ogx_open_client 1.0.3.dev0__py3-none-any.whl

ogx_client/lib/cli/inference/inference.py

@@ -0,0 +1,88 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import traceback
+
+import click
+from rich.console import Console
+
+from ...inference.event_logger import EventLogger
+from ..common.utils import handle_client_errors
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def inference():
+    """Inference (chat)."""
+
+
+@click.command("chat-completion")
+@click.help_option("-h", "--help")
+@click.option("--message", help="Message")
+@click.option("--stream", is_flag=True, help="Streaming", default=False)
+@click.option("--session", is_flag=True, help="Start a Chat Session", default=False)
+@click.option("--model-id", required=False, help="Model ID")
+@click.pass_context
+@handle_client_errors("inference chat-completion")
+def chat_completion(ctx, message: str, stream: bool, session: bool, model_id: str | None):
+    """Show available inference chat completion endpoints on distribution endpoint"""
+    if not message and not session:
+        click.secho(
+            "you must specify either --message or --session",
+            fg="red",
+        )
+        raise click.exceptions.Exit(1)
+    client = ctx.obj["client"]
+    console = Console()
+
+    if not model_id:
+        available_models = [model.identifier for model in client.models.list() if model.model_type == "llm"]
+        model_id = available_models[0]
+
+    messages = []
+    if message:
+        messages.append({"role": "user", "content": message})
+        response = client.chat.completions.create(
+            model=model_id,
+            messages=messages,
+            stream=stream,
+        )
+        if not stream:
+            console.print(response)
+        else:
+            for event in EventLogger().log(response):
+                event.print()
+    if session:
+        chat_session(client=client, model_id=model_id, messages=messages, console=console)
+
+
+def chat_session(client, model_id: str | None, messages: list[dict[str, str]], console: Console):
+    """Run an interactive chat session with the served model"""
+    while True:
+        try:
+            message = input(">>> ")
+            if message in ["\\q", "quit"]:
+                console.print("Exiting")
+                break
+            messages.append({"role": "user", "content": message})
+            response = client.chat.completions.create(
+                model=model_id,
+                messages=messages,
+                stream=True,
+            )
+            for event in EventLogger().log(response):
+                event.print()
+        except Exception as exc:
+            traceback.print_exc()
+            console.print(f"Error in chat session {exc}")
+            break
+        except KeyboardInterrupt:
+            console.print("\nDetected user interrupt, exiting")
+            break
+
+
+# Register subcommands
+inference.add_command(chat_completion)

ogx_client/lib/cli/inspect/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .inspect import inspect
+
+__all__ = ["inspect"]

ogx_client/lib/cli/inspect/inspect.py

@@ -0,0 +1,19 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+
+from .version import inspect_version
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def inspect():
+    """Inspect server configuration."""
+
+
+# Register subcommands
+inspect.add_command(inspect_version)

ogx_client/lib/cli/inspect/version.py

@@ -0,0 +1,22 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+from rich.console import Console
+
+from ..common.utils import handle_client_errors
+
+
+@click.command("version")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("inspect version")
+def inspect_version(ctx):
+    """Show OGX version on distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+    version_response = client.inspect.version()
+    console.print(version_response)

ogx_client/lib/cli/models/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .models import models
+
+__all__ = ["models"]

ogx_client/lib/cli/models/models.py

@@ -0,0 +1,149 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from ..common.utils import handle_client_errors
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def models():
+    """Manage GenAI models."""
+
+
+@click.command(name="list", help="Show available llama models at distribution endpoint")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("list models")
+def list_models(ctx):
+    client = ctx.obj["client"]
+    console = Console()
+
+    response = client.models.list()
+    if response:
+        table = Table(
+            show_lines=True,  # Add lines between rows for better readability
+            padding=(0, 1),  # Add horizontal padding
+            expand=True,  # Allow table to use full width
+        )
+
+        # Configure columns with specific styling
+        table.add_column("model_type", style="blue")
+        table.add_column("identifier", style="bold cyan", no_wrap=True, overflow="fold")
+        table.add_column("provider_resource_id", style="yellow", no_wrap=True, overflow="fold")
+        table.add_column("metadata", style="magenta", max_width=30, overflow="fold")
+        table.add_column("provider_id", style="green", max_width=20)
+
+        for item in response:
+            table.add_row(
+                item.model_type,
+                item.identifier,
+                item.provider_resource_id,
+                str(item.metadata or ""),
+                item.provider_id,
+            )
+
+        # Create a title for the table
+        console.print("\n[bold]Available Models[/bold]\n")
+        console.print(table)
+        console.print(f"\nTotal models: {len(response)}\n")
+
+
+@click.command(name="get")
+@click.help_option("-h", "--help")
+@click.argument("model_id")
+@click.pass_context
+@handle_client_errors("get model details")
+def get_model(ctx, model_id: str):
+    """Show details of a specific model at the distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    models_get_response = client.models.retrieve(model_id=model_id)
+
+    if not models_get_response:
+        console.print(
+            f"Model {model_id} is not found at distribution endpoint. "
+            "Please ensure endpoint is serving specified model.",
+            style="bold red",
+        )
+        return
+
+    headers = sorted(models_get_response.__dict__.keys())
+    table = Table()
+    for header in headers:
+        table.add_column(header)
+
+    table.add_row(*[str(models_get_response.__dict__[header]) for header in headers])
+    console.print(table)
+
+
+class JSONParamType(click.ParamType):
+    name = "json"
+
+    def convert(self, value, param, ctx):
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError as e:
+            self.fail(f"Invalid JSON: {e}", param, ctx)
+
+
+@click.command(name="register", help="Register a new model at distribution endpoint")
+@click.help_option("-h", "--help")
+@click.argument("model_id")
+@click.option("--provider-id", help="Provider ID for the model", default=None)
+@click.option("--provider-model-id", help="Provider's model ID", default=None)
+@click.option("--metadata", type=JSONParamType(), help="JSON metadata for the model", default=None)
+@click.option("--model-type", type=click.Choice(["llm", "embedding"]), default="llm", help="Model type: llm, embedding")
+@click.pass_context
+@handle_client_errors("register model")
+def register_model(
+    ctx,
+    model_id: str,
+    provider_id: str | None,
+    provider_model_id: str | None,
+    metadata: str | None,
+    model_type: str | None,
+):
+    """Register a new model at distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    response = client.models.register(
+        model_id=model_id,
+        provider_id=provider_id,
+        provider_model_id=provider_model_id,
+        metadata=metadata,
+        model_type=model_type,
+    )
+    if response:
+        console.print(f"[green]Successfully registered model {model_id}[/green]")
+
+
+@click.command(name="unregister", help="Unregister a model from distribution endpoint")
+@click.help_option("-h", "--help")
+@click.argument("model_id")
+@click.pass_context
+@handle_client_errors("unregister model")
+def unregister_model(ctx, model_id: str):
+    client = ctx.obj["client"]
+    console = Console()
+
+    response = client.models.unregister(model_id=model_id)
+    if response:
+        console.print(f"[green]Successfully deleted model {model_id}[/green]")
+
+
+# Register subcommands
+models.add_command(list_models)
+models.add_command(get_model)
+models.add_command(register_model)
+models.add_command(unregister_model)

ogx_client/lib/cli/ogx_client.py

@@ -0,0 +1,111 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+from importlib.metadata import version
+
+import click
+import yaml
+
+from ... import OgxClient
+
+
+def _get_version():
+    for dist_name in ("ogx-client", "ogx-open-client"):
+        try:
+            return version(dist_name)
+        except Exception:
+            continue
+    return "unknown"
+
+
+from .configure import configure
+from .constants import get_config_file_path
+from .datasets import datasets
+from .eval import eval
+from .eval_tasks import eval_tasks
+from .inference import inference
+from .inspect import inspect
+from .models import models
+from .post_training import post_training
+from .providers import providers
+from .scoring_functions import scoring_functions
+from .toolgroups import toolgroups
+from .vector_stores import vector_stores
+
+
+@click.group()
+@click.help_option("-h", "--help")
+@click.version_option(version=_get_version(), prog_name="ogx-client")
+@click.option("--endpoint", type=str, help="OGX distribution endpoint", default="")
+@click.option("--api-key", type=str, help="OGX distribution API key", default="")
+@click.option("--config", type=str, help="Path to config file", default=None)
+@click.pass_context
+def ogx_client(ctx, endpoint: str, api_key: str, config: str | None):
+    """Welcome to the ogx-client CLI - a command-line interface for interacting with OGX"""
+    ctx.ensure_object(dict)
+
+    # If no config provided, check default location
+    if config and endpoint:
+        raise ValueError("Cannot use both config and endpoint")
+
+    if config is None:
+        default_config = get_config_file_path()
+        if default_config.exists():
+            config = str(default_config)
+
+    if config:
+        try:
+            with open(config) as f:
+                config_dict = yaml.safe_load(f)
+                endpoint = config_dict.get("endpoint", endpoint)
+                api_key = config_dict.get("api_key", "")
+        except Exception as e:
+            click.echo(f"Error loading config from {config}: {str(e)}", err=True)
+            click.echo("Falling back to HTTP client with endpoint", err=True)
+
+    if endpoint == "":
+        endpoint = "http://localhost:8321"
+
+    default_headers = {}
+    if api_key != "":
+        default_headers = {
+            "Authorization": f"Bearer {api_key}",
+        }
+
+    client = OgxClient(
+        base_url=endpoint,
+        provider_data={
+            "fireworks_api_key": os.environ.get("FIREWORKS_API_KEY", ""),
+            "together_api_key": os.environ.get("TOGETHER_API_KEY", ""),
+            "openai_api_key": os.environ.get("OPENAI_API_KEY", ""),
+        },
+        default_headers=default_headers,
+    )
+    ctx.obj = {"client": client}
+
+
+# Register all subcommands
+ogx_client.add_command(models, "models")
+ogx_client.add_command(vector_stores, "vector_stores")
+ogx_client.add_command(eval_tasks, "eval_tasks")
+ogx_client.add_command(providers, "providers")
+ogx_client.add_command(datasets, "datasets")
+ogx_client.add_command(configure, "configure")
+ogx_client.add_command(scoring_functions, "scoring_functions")
+ogx_client.add_command(eval, "eval")
+ogx_client.add_command(inference, "inference")
+ogx_client.add_command(post_training, "post_training")
+ogx_client.add_command(inspect, "inspect")
+ogx_client.add_command(toolgroups, "toolgroups")
+
+
+def main():
+    ogx_client()
+
+
+if __name__ == "__main__":
+    main()

ogx_client/lib/cli/post_training/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .post_training import post_training
+
+__all__ = ["post_training"]

ogx_client/lib/cli/post_training/post_training.py

@@ -0,0 +1,114 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+
+import click
+from rich.console import Console
+
+from ....models.algorithm_config import AlgorithmConfig
+from ....models.training_config import TrainingConfig
+from ..common.utils import handle_client_errors
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def post_training():
+    """Post-training."""
+
+
+@click.command("supervised_fine_tune")
+@click.help_option("-h", "--help")
+@click.option("--job-uuid", required=True, help="Job UUID")
+@click.option("--model", required=True, help="Model ID")
+@click.option("--algorithm-config", required=True, help="Algorithm Config")
+@click.option("--training-config", required=True, help="Training Config")
+@click.option("--checkpoint-dir", required=False, help="Checkpoint Config", default=None)
+@click.pass_context
+@handle_client_errors("post_training supervised_fine_tune")
+def supervised_fine_tune(
+    ctx,
+    job_uuid: str,
+    model: str,
+    algorithm_config: AlgorithmConfig,
+    training_config: TrainingConfig,
+    checkpoint_dir: str | None,
+):
+    """Kick off a supervised fine tune job"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    post_training_job = client.post_training.supervised_fine_tune(
+        job_uuid=job_uuid,
+        model=model,
+        algorithm_config=algorithm_config,
+        training_config=training_config,
+        checkpoint_dir=checkpoint_dir,
+        # logger_config and hyperparam_search_config haven't been used yet
+        logger_config={},
+        hyperparam_search_config={},
+    )
+    console.print(post_training_job.job_uuid)
+
+
+@click.command("list")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("post_training get_training_jobs")
+def get_training_jobs(ctx):
+    """Show the list of available post training jobs"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    post_training_jobs = client.post_training.job.list()
+    console.print([post_training_job.job_uuid for post_training_job in post_training_jobs])
+
+
+@click.command("status")
+@click.help_option("-h", "--help")
+@click.option("--job-uuid", required=True, help="Job UUID")
+@click.pass_context
+@handle_client_errors("post_training get_training_job_status")
+def get_training_job_status(ctx, job_uuid: str):
+    """Show the status of a specific post training job"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    job_status_reponse = client.post_training.job.status(job_uuid=job_uuid)
+    console.print(job_status_reponse)
+
+
+@click.command("artifacts")
+@click.help_option("-h", "--help")
+@click.option("--job-uuid", required=True, help="Job UUID")
+@click.pass_context
+@handle_client_errors("post_training get_training_job_artifacts")
+def get_training_job_artifacts(ctx, job_uuid: str):
+    """Get the training artifacts of a specific post training job"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    job_artifacts = client.post_training.job.artifacts(job_uuid=job_uuid)
+    console.print(job_artifacts)
+
+
+@click.command("cancel")
+@click.help_option("-h", "--help")
+@click.option("--job-uuid", required=True, help="Job UUID")
+@click.pass_context
+@handle_client_errors("post_training cancel_training_job")
+def cancel_training_job(ctx, job_uuid: str):
+    """Cancel the training job"""
+    client = ctx.obj["client"]
+
+    client.post_training.job.cancel(job_uuid=job_uuid)
+
+
+# Register subcommands
+post_training.add_command(supervised_fine_tune)
+post_training.add_command(get_training_jobs)
+post_training.add_command(get_training_job_status)
+post_training.add_command(get_training_job_artifacts)
+post_training.add_command(cancel_training_job)

ogx_client/lib/cli/providers/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .providers import providers
+
+__all__ = ["providers"]

ogx_client/lib/cli/providers/inspect.py

@@ -0,0 +1,33 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+import yaml
+from rich.console import Console
+
+from ..common.utils import handle_client_errors
+
+
+@click.command(name="inspect")
+@click.argument("provider_id")
+@click.pass_context
+@handle_client_errors("inspect providers")
+def inspect_provider(ctx, provider_id):
+    """Show specific provider configuration on distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    providers_response = client.providers.retrieve(provider_id=provider_id)
+
+    if not providers_response:
+        click.secho("Provider not found", fg="red")
+        raise click.exceptions.Exit(1)
+
+    console.print(f"provider_id={providers_response.provider_id}")
+    console.print(f"provider_type={providers_response.provider_type}")
+    console.print("config:")
+    for line in yaml.dump(providers_response.config, indent=2).split("\n"):
+        console.print(line)

ogx_client/lib/cli/providers/list.py

@@ -0,0 +1,32 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from ..common.utils import handle_client_errors
+
+
+@click.command("list")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("list providers")
+def list_providers(ctx):
+    """Show available providers on distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+    headers = ["API", "Provider ID", "Provider Type"]
+
+    providers_response = client.providers.list()
+    table = Table()
+    for header in headers:
+        table.add_column(header)
+
+    for response in providers_response:
+        table.add_row(response.api, response.provider_id, response.provider_type)
+
+    console.print(table)

ogx_client/lib/cli/providers/providers.py

@@ -0,0 +1,21 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+
+from .inspect import inspect_provider
+from .list import list_providers
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def providers():
+    """Manage API providers."""
+
+
+# Register subcommands
+providers.add_command(list_providers)
+providers.add_command(inspect_provider)

ogx_client/lib/cli/scoring_functions/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .scoring_functions import scoring_functions
+
+__all__ = ["scoring_functions"]

ogx_client/lib/cli/scoring_functions/list.py

@@ -0,0 +1,38 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import click
+from rich.console import Console
+from rich.table import Table
+
+from ..common.utils import handle_client_errors
+
+
+@click.command("list")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("list scoring functions")
+def list_scoring_functions(ctx):
+    """Show available scoring functions on distribution endpoint"""
+
+    client = ctx.obj["client"]
+    console = Console()
+    headers = [
+        "identifier",
+        "provider_id",
+        "description",
+        "type",
+    ]
+
+    scoring_functions_list_response = client.scoring_functions.list()
+    if scoring_functions_list_response:
+        table = Table()
+        for header in headers:
+            table.add_column(header)
+
+        for item in scoring_functions_list_response:
+            table.add_row(*[str(getattr(item, header)) for header in headers])
+        console.print(table)