diff --git a/src/pieces/commands/list_command.py b/src/pieces/commands/list_command.py index 86c565a..b7ef9e8 100644 --- a/src/pieces/commands/list_command.py +++ b/src/pieces/commands/list_command.py @@ -1,7 +1,9 @@ from collections.abc import Iterable +import threading from pieces.settings import Settings from pieces.utils import PiecesSelectMenu +from pieces.wrapper.basic_identifier.asset import BasicAsset from .change_model import change_model from .assets_command import check_assets_existence, AssetsCommands @@ -25,20 +27,24 @@ def list_command(cls, **kwargs): @classmethod @check_assets_existence - def list_assets(cls, max_assets: int = 10,**kwargs): - assets = kwargs.get("assets",Settings.pieces_client.assets()[:max_assets]) - - select_menu = PiecesSelectMenu( - [ - (f"{i}: {asset.name}", {"asset_id":asset.id,**kwargs}) - for i, asset in enumerate(assets,start=1) - ], - AssetsCommands.open_asset,kwargs.get("footer")) + def list_assets(cls, **kwargs): + assets = kwargs.get("assets",[BasicAsset(item.id) for item in BasicAsset.get_identifiers()]) + + select_menu = PiecesSelectMenu([], AssetsCommands.open_asset,kwargs.get("footer")) + def update_assets(): + for i,asset in enumerate(assets,start=1): + select_menu.add_entry( + (f"{i}: {asset.name}", {"asset_id":asset.id,**kwargs})) + + threading.Thread(target=update_assets).start() select_menu.run() + + + @classmethod def list_models(cls): - models = [(f"{idx}: {model_name}", {"MODEL_INDEX":idx,"show_warning":False}) + models = [(f"{idx}: {model_name}", {"MODEL_INDEX":idx}) for idx, model_name in enumerate(Settings.pieces_client.available_models_names, start=1)] select_menu = PiecesSelectMenu(models, change_model,f"Currently using: {Settings.get_model()}") select_menu.run() diff --git a/src/pieces/utils.py b/src/pieces/utils.py index f96ef15..2948f01 100644 --- a/src/pieces/utils.py +++ b/src/pieces/utils.py @@ -11,7 +11,7 @@ from prompt_toolkit.layout.containers import HSplit, Window from prompt_toolkit.layout.controls import FormattedTextControl from prompt_toolkit.styles import Style -from typing import List, Tuple, Callable,Optional +from typing import Any, List, Tuple, Callable,Optional # Used to create a valid file name when opening to "Opened Snippets" def sanitize_filename(name): @@ -54,11 +54,9 @@ def get_file_extension(language): # Return the corresponding file extension or default to '.txt' if not found return extension_mapping.get(language, '.txt') - - class PiecesSelectMenu: def __init__(self, menu_options: List[Tuple], on_enter_callback: Callable, footer_text: Optional[str] = None): - self.menu_options = menu_options + self.menu_options = list(menu_options) # Ensure it's a list self.on_enter_callback = on_enter_callback self.current_selection = 0 self.footer_text = footer_text @@ -68,6 +66,12 @@ def update_visible_range(self): terminal_size = shutil.get_terminal_size() self.visible_start = 0 self.visible_end = terminal_size.lines - 2 + + def add_entry(self, entry: Tuple[str, Any]): + """Add a new entry to the menu and update the layout.""" + self.menu_options.append(entry) + self.update_visible_range() + self.update_app() def get_menu_text(self): result = [] @@ -77,6 +81,15 @@ def get_menu_text(self): else: result.append(('class:unselected', f' {option[0]}\n')) return result + + def update_app(self): + if hasattr(self, "app"): + self.menu_window.content = FormattedTextControl(text=self.get_menu_text) + self.app.invalidate() + # self.app.layout.focus(self.menu_window) + else: + pass + # raise ValueError("App not initialized") def run(self): bindings = KeyBindings() @@ -104,8 +117,7 @@ def select_option(event): args = self.menu_options[self.current_selection][1] event.app.exit(result=args) - menu_control = FormattedTextControl(text=self.get_menu_text) - self.menu_window = Window(content=menu_control, always_hide_cursor=True) + self.menu_window = Window(content=FormattedTextControl(text=self.get_menu_text), always_hide_cursor=True) layout_items = [self.menu_window] @@ -121,8 +133,8 @@ def select_option(event): 'unselected': '' }) - app = Application(layout=layout, key_bindings=bindings, style=style, full_screen=True) - args = app.run() + self.app = Application(layout=layout, key_bindings=bindings, style=style, full_screen=True) + args = self.app.run() if isinstance(args, list): self.on_enter_callback(*args) @@ -130,4 +142,3 @@ def select_option(event): self.on_enter_callback(args) elif isinstance(args, dict): self.on_enter_callback(**args) - diff --git a/src/pieces/wrapper/basic_identifier/asset.py b/src/pieces/wrapper/basic_identifier/asset.py index b6c26ff..ae7bf6a 100644 --- a/src/pieces/wrapper/basic_identifier/asset.py +++ b/src/pieces/wrapper/basic_identifier/asset.py @@ -21,19 +21,23 @@ class BasicAsset(Basic): """ A wrapper class for managing assets. """ - @staticmethod - def identifiers_snapshot(): + @classmethod + def identifiers_snapshot(cls): if AssetSnapshot.identifiers_snapshot: return AssetSnapshot.identifiers_snapshot + + AssetSnapshot.identifiers_snapshot = {item.id:None for item in cls.get_identifiers()} + return AssetSnapshot.identifiers_snapshot + + @staticmethod + def get_identifiers(): + """ + :returns: The assets id + """ assets_api = AssetSnapshot.pieces_client.assets_api - # Call the API to get assets identifiers api_response = assets_api.assets_identifiers_snapshot() - - # Extract the 'id' values from each item in the 'iterable' list - AssetSnapshot.identifiers_snapshot = {item.id:None for item in api_response.iterable} - - return AssetSnapshot.identifiers_snapshot + return api_response.iterable @property def asset(self) -> "Asset":