adding preprompt concept

Author: blainekasten

src/together/lib/cli/__init__.py

@@ -9,11 +9,12 @@
 import httpx
 from cyclopts import App, MissingArgumentError, Parameter
 
+from together.lib.cli.logger.config import CLIConfig
 from together import AsyncTogether
 from together._exceptions import APIError
 from together._version import __version__
 from together._utils._logs import setup_logging
-from together.lib.cli.logger.prompt import PromptParameter
+from together.lib.cli.logger.prompt import PromptParameter, console
 
 app = App(
     name="together",
@@ -25,16 +26,6 @@
 app['--version'].group = "Parameters"
 app['--help'].group = "Parameters"
 
-class Config:
-    client: AsyncTogether
-    non_interactive: bool
-    json: bool
-
-    def __init__(self, client: AsyncTogether, non_interactive: bool, json: bool):
-        self.client = client
-        self.non_interactive = non_interactive
-        self.json = json
-
 def _create_client(
     api_key: Optional[str],
     base_url: Optional[str],
@@ -86,7 +77,7 @@ async def _launcher(
         os.environ.setdefault("TOGETHER_LOG", "debug")
         setup_logging()
     client = _create_client(api_key, base_url, timeout, max_retries)
-    config = Config(
+    config = CLIConfig(
         client=client,
         non_interactive=non_interactive or False,
         json=json or False,
@@ -104,7 +95,6 @@ async def run_command():
                     remaining.append(value)
 
             kwargs = dict(bound.kwargs)
-            kwargs["config"] = config
             if "config" in extra:
                 kwargs["config"] = config
             result = command(*bound.args, **kwargs)
@@ -130,12 +120,17 @@ async def run_command():
 
             value: str | None = None
             if prompt is not None:
+                await prompt.preprompt(config)
                 value = await prompt.prompt(e.argument.name)
                 print("") # Push a blank line for nicer output
                 remaining.append(e.argument.name)
                 remaining.append(value)
                 await run_command()
-        except (KeyboardInterrupt, SystemExit):
+            else:
+                # TODO: Better design this
+                print("Missing required argument", e.argument.name)
+                sys.exit(1)
+        except KeyboardInterrupt:
             pass
         except APIError as e:
             error_msg = ""
@@ -146,12 +141,16 @@ async def run_command():
             print(f"Failed", file=sys.stderr)
             print(f"{error_msg}", file=sys.stderr)
             sys.exit(1)
-        except Exception as e:
-            print(f"Failed", file=sys.stderr)
-            print(f"An unexpected error occurred - {e!s}", file=sys.stderr)
-            sys.exit(1)
+        # except Exception as e:
+            # print(f"Failed", file=sys.stderr)
+            # print(f"An unexpected error occurred - {e!s}", file=sys.stderr)
+            # sys.exit(1)
     try:
         await run_command()
+    # except Exception as e:
+    #     print(f"Failed", file=sys.stderr)
+    #     print(f"An unexpected error occurred - {e!s}", file=sys.stderr)
+    #     sys.exit(1)
     finally:
         await client.close()
 

src/together/lib/cli/api/endpoints/create.py

@@ -6,14 +6,12 @@
 
 from cyclopts import Parameter
 
-from together import APIError, AsyncTogether, omit
+from together import APIError, omit
+from together.lib.cli.logger.config import CLIConfig
 
-from together.lib.cli.api.endpoints._utils import handle_endpoint_api_errors
 
 from .hardware import hardware as list_hardware
 
-
-@handle_endpoint_api_errors("Endpoints")
 async def create(
     model: str,
     min_replicas: int = 1,
@@ -26,9 +24,8 @@ async def create(
     inactive_timeout: Optional[int] = None,
     availability_zone: Optional[str] = None,
     wait: bool = False,
-    json_output: bool = False,
     *,
-    client: Annotated[AsyncTogether, Parameter(parse=False)],
+    config: Annotated[CLIConfig, Parameter(parse=False)],
 ) -> None:
     """Create a new dedicated inference endpoint."""
     if min_replicas > max_replicas:
@@ -40,7 +37,7 @@ async def create(
 
     if availability_zone:
         try:
-            valid_zones = await client.endpoints.list_avzones()
+            valid_zones = await config.client.endpoints.list_avzones()
             if availability_zone not in valid_zones.avzones:
                 print(f"Error: Invalid availability zone '{availability_zone}'", file=sys.stderr)
                 if valid_zones.avzones:
@@ -51,19 +48,19 @@ async def create(
         except Exception:
             pass
 
-    if json_output and wait:
+    if config.json and wait:
         print("Error: --json and --wait cannot be used together.", file=sys.stderr)
         return
 
-    if no_prompt_cache is not None and not json_output:
+    if no_prompt_cache is not None and not config.json:
         print("Warning: --no-prompt-cache is deprecated and no longer has any effect.", file=sys.stderr)
 
     if hardware is None:
         print("Error: --hardware is required", file=sys.stderr)
         sys.exit(1)
 
     try:
-        response = await client.endpoints.create(
+        response = await config.client.endpoints.create(
             model=model,
             hardware=hardware,
             autoscaling={"min_replicas": min_replicas, "max_replicas": max_replicas},
@@ -74,7 +71,7 @@ async def create(
             extra_query={"availability_zone": availability_zone or omit},
         )
     except APIError as e:
-        if json_output:
+        if config.json:
             raise e
         error_msg = str(e.args[0]).lower() if e.args else ""
         if (
@@ -85,7 +82,7 @@ async def create(
         ):
             print("Invalid hardware selected.", file=sys.stderr)
             print("\nAvailable hardware options:", file=sys.stderr)
-            await list_hardware(model=model, json_output=False, available=True, client=client)
+            await list_hardware(model=model, json_output=False, available=True, client=config.client)
             sys.exit(1)
         elif "model" in error_msg and (
             "not found" in error_msg
@@ -99,7 +96,7 @@ async def create(
             sys.exit(1)
         raise e
 
-    if json_output:
+    if config.json:
         print(response.model_dump_json(indent=2))
         return
 
@@ -122,7 +119,7 @@ async def create(
 
     if wait:
         print("Waiting for endpoint to be ready...", file=sys.stderr)
-        while (await client.endpoints.retrieve(response.id)).state != "STARTED":
+        while (await config.client.endpoints.retrieve(response.id)).state != "STARTED":
             await asyncio.sleep(1)
         print("Endpoint ready", file=sys.stderr)
     print(response.id)

src/together/lib/cli/api/endpoints/start.py

@@ -3,36 +3,54 @@
 import json as json_lib
 import sys
 from typing import Annotated
+from typing_extensions import override
 
 from cyclopts import Parameter
 
 import asyncio
+import questionary
 
-from together import AsyncTogether
-
-from together.lib.cli.api.endpoints._utils import handle_endpoint_api_errors
+from together.lib.cli.logger.config import CLIConfig
+from together.lib.cli.logger.prompt import PromptParameter, console
 from together.lib.utils.serializer import datetime_serializer
 
+class LoadEndpointsPrompt(PromptParameter):
+    @override
+    async def preprompt(self, config: CLIConfig):
+        with console.status(
+            "[progress.description]Loading endpoints...[/progress.description]",
+            spinner="dots",
+            spinner_style="bar.pulse",
+        ):
+            endpoints = await config.client.endpoints.list()
+            self.choices = []
+            for endpoint in endpoints.data:
+                # This shouldn't happen.. but does happen sometimes...
+                if endpoint.id is None: # type: ignore
+                    continue
+
+                if endpoint.state != "STARTED":
+                    self.choices.append(questionary.Choice(title=[("", endpoint.name), ("class:disabled", " ({})".format(endpoint.id))], value=endpoint.id))
+
 
-@handle_endpoint_api_errors("Endpoints")
 async def start(
-    endpoint_id: str,
+    endpoint_id: Annotated[str, Parameter(required=True, help="The ID of the endpoint to start"), LoadEndpointsPrompt(message="Enter the endpoint ID")],
     wait: bool = False,
-    json_output: bool = False,
     *,
-    client: Annotated[AsyncTogether, Parameter(parse=False)],
+    config: Annotated[CLIConfig, Parameter(parse=False)],
 ) -> None:
     """Start a dedicated inference endpoint."""
-    response = await client.endpoints.update(endpoint_id, state="STARTED")
+    response = await config.client.endpoints.update(endpoint_id, state="STARTED")
+    
 
-    if json_output:
+    if config.json:
         print(json_lib.dumps(response.model_dump(), default=datetime_serializer, indent=2))
         return
 
     print("Successfully marked endpoint as starting", file=sys.stderr)
     if wait:
         print("Waiting for endpoint to start...", file=sys.stderr)
-        while (await client.endpoints.retrieve(endpoint_id)).state != "STARTED":
+        while (await config.client.endpoints.retrieve(endpoint_id)).state != "STARTED":
             await asyncio.sleep(1)
         print("Endpoint started", file=sys.stderr)
     print(endpoint_id)

src/together/lib/cli/api/models/list.py

@@ -5,24 +5,23 @@
 from tabulate import tabulate
 
 from cyclopts import Parameter
+from rich import print_json
 
-from together import AsyncTogether, omit
-from together.lib.utils.serializer import datetime_serializer
+
+from together import omit
+from together.lib.cli.logger.config import CLIConfig
+from together._utils._json import openapi_dumps
 
 async def list_(
     type_: Annotated[Optional[Literal["dedicated"]], Parameter(name="type")] = None,
-    json_output: bool = False,
     *,
-    client: Annotated[AsyncTogether, Parameter(parse=False)],
+    config: Annotated[CLIConfig, Parameter(parse=False)],
 ) -> None:
     """List models."""
-    models_list = await client.models.list(dedicated=type_ == "dedicated" if type_ else omit)
-
-    if json_output:
-        import json as json_lib
+    models_list = await config.client.models.list(dedicated=type_ == "dedicated" if type_ else omit)
 
-        items = [model.model_dump() for model in models_list]
-        print(json_lib.dumps(items, indent=2, default=datetime_serializer))
+    if config.json:
+        print_json(openapi_dumps(models_list).decode())
         return
 
     display_list: List[Dict[str, Any]] = []

src/together/lib/cli/api/models/upload.py

@@ -8,7 +8,7 @@
 
 from together import TogetherError, omit
 
-from together.lib.cli import Config
+from together.lib.cli.logger.config import CLIConfig
 from together.lib.cli.logger.prompt import PromptParameter
 from together.types.model_upload_response import ModelUploadResponse
 from together._utils._json import openapi_dumps
@@ -56,7 +56,7 @@ async def upload(
     base_model: Optional[str] = None,
     lora_model: Optional[str] = None,
     *,
-    config: Annotated[Config, Parameter(parse=False)],
+    config: Annotated[CLIConfig, Parameter(parse=False)],
 ) -> None:
     """Upload a custom model or adapter from Hugging Face or S3."""
     response: ModelUploadResponse = await config.client.models.upload(

src/together/lib/cli/logger/config.py

@@ -0,0 +1,11 @@
+from together import AsyncTogether
+
+class CLIConfig:
+    client: AsyncTogether
+    non_interactive: bool
+    json: bool
+
+    def __init__(self, client: AsyncTogether, non_interactive: bool, json: bool):
+        self.client = client
+        self.non_interactive = non_interactive
+        self.json = json

src/together/lib/cli/logger/prompt.py

@@ -2,6 +2,9 @@
 from rich import print
 
 import questionary
+from together.lib.cli.logger.config import CLIConfig
+from rich.theme import Theme
+from rich.console import Console
 
 custom_style_fancy = questionary.Style([
     ('qmark', 'fg:#caaef5 bold'),       # token in front of the question
@@ -16,6 +19,41 @@
     ('disabled', 'fg:#858585 italic')   # disabled choices for select and checkbox prompts
 ])
 
+custom_theme = Theme({
+    # Text styles
+    "primary": "#caaef5",       # Purple 300 ⭐ (lighter when bold)
+    "secondary": "dim #caaef5",          # Purple 500 ⭐ (mid-tone without bold)
+    "accent": "#ff68d4",             # Pink 500 ⭐
+    "muted": "#98a0b3",              # Grey 400 ⭐
+    "dim": "dim #626b84",            # Grey 600 ⭐
+
+    # Semantic styles
+    "success": "bold #0dce74",       # Green 400 ⭐
+    "info": "#64afff",               # Blue 500 ⭐
+    "warning": "bold #ff815d",       # Red 500 ⭐
+    "error": "bold #c63800",         # Red 700 ⭐
+
+    # UI elements
+    "prompt": "#ba92ff",             # Purple 500 ⭐ (no bold)
+    "prompt.choices": "#caaef5",     # Purple 300 ⭐
+    "prompt.default": "dim #98a0b3", # Grey 400 ⭐
+
+    # Table styles
+    "table.header": "#414858",  # Purple 300 ⭐ (lighter when bold)
+    "table.border": "#626b84",       # Grey 600 ⭐
+    "table.row": "#c4c9d4",          # Grey 300 ⭐
+
+    # Progress/Loading
+    "progress.description": "#caaef5",     # Purple 300 ⭐
+    "progress.percentage": "bold #caaef5", # Purple 300 ⭐ (lighter when bold)
+    "bar.complete": "#ba92ff",             # Purple 500 ⭐ (no bold)
+    "bar.finished": "#0dce74",             # Green 400 ⭐
+    "bar.pulse": "#ff68d4",                # Pink 500 ⭐
+})
+
+console = Console(theme=custom_theme)
+
+
 class NameValidator(questionary.Validator):
     def validate(self, document):
         if document.text.count(" ") > 0:
@@ -27,18 +65,21 @@ def validate(self, document):
 class PromptParameter:
     message: str | None = None
     instructions: str | None = None
-    choices: list[str] | None = None
+    choices: list[str | questionary.Choice] | None = None
 
-    def __init__(self, message: str | None = None, instructions: str | None = None, choices: list[str] | None = None):
+    def __init__(self, message: str | None = None, instructions: str | None = None, choices: list[str | questionary.Choice] | None = None):
         self.message = message
         self.instructions = instructions
         self.choices = choices
 
+    async def preprompt(self, _config: CLIConfig):
+        pass
+
     async def prompt(self, field: str) -> str:
         if self.instructions is not None:
             print(f"[dim]{self.instructions}[/dim]")
-        
+
         if self.choices is not None:
-            return await questionary.select(self.message or field, choices=self.choices, style=custom_style_fancy).unsafe_ask_async()
+            return await questionary.select(self.message or field, choices=self.choices, style=custom_style_fancy, show_selected=True).unsafe_ask_async()
 
         return await questionary.text(self.message or field, instruction="\n→", style=custom_style_fancy, validate=NameValidator).unsafe_ask_async()