ogx_open_client 1.0.3.dev0__py3-none-any.whl

ogx_open_client/lib/cli/scoring_functions/scoring_functions.py

@@ -0,0 +1,62 @@
+# 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
+import yaml
+
+from .list import list_scoring_functions
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def scoring_functions():
+    """Manage scoring functions."""
+
+
+@scoring_functions.command()
+@click.help_option("-h", "--help")
+@click.option("--scoring-fn-id", required=True, help="Id of the scoring function")
+@click.option("--description", required=True, help="Description of the scoring function")
+@click.option("--return-type", type=str, required=True, help="Return type of the scoring function")
+@click.option("--provider-id", type=str, help="Provider ID for the scoring function", default=None)
+@click.option("--provider-scoring-fn-id", type=str, help="Provider's scoring function ID", default=None)
+@click.option("--params", type=str, help="Parameters for the scoring function in JSON format", default=None)
+@click.pass_context
+def register(
+    ctx,
+    scoring_fn_id: str,
+    description: str,
+    return_type: str,
+    provider_id: str | None,
+    provider_scoring_fn_id: str | None,
+    params: str | None,
+):
+    """Register a new scoring function"""
+    client = ctx.obj["client"]
+
+    if params:
+        try:
+            params = json.loads(params)
+        except json.JSONDecodeError as err:
+            raise click.BadParameter("Parameters must be valid JSON") from err
+
+    response = client.scoring_functions.register(
+        scoring_fn_id=scoring_fn_id,
+        description=description,
+        return_type=json.loads(return_type),
+        provider_id=provider_id,
+        provider_scoring_fn_id=provider_scoring_fn_id,
+        params=params,
+    )
+    if response:
+        click.echo(yaml.dump(response.dict()))
+
+
+# Register subcommands
+scoring_functions.add_command(list_scoring_functions)
+scoring_functions.add_command(register)

ogx_open_client/lib/cli/toolgroups/__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 .toolgroups import toolgroups
+
+__all__ = ["toolgroups"]

ogx_open_client/lib/cli/toolgroups/toolgroups.py

@@ -0,0 +1,129 @@
+# 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 ...._types import NOT_GIVEN, NotGiven
+from ....types import toolgroup_register_params
+from ..common.utils import handle_client_errors
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def toolgroups():
+    """Manage available tool groups."""
+
+
+@click.command(name="list", help="Show available toolgroups at distribution endpoint")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("list toolgroups")
+def list_toolgroups(ctx):
+    client = ctx.obj["client"]
+    console = Console()
+
+    headers = ["identifier", "provider_id", "args", "mcp_endpoint"]
+    response = client.toolgroups.list()
+    if response:
+        table = Table()
+        for header in headers:
+            table.add_column(header)
+
+        for item in response:
+            row = [str(getattr(item, header)) for header in headers]
+            table.add_row(*row)
+        console.print(table)
+
+
+@click.command(name="get")
+@click.help_option("-h", "--help")
+@click.argument("toolgroup_id")
+@click.pass_context
+@handle_client_errors("get toolgroup details")
+def get_toolgroup(ctx, toolgroup_id: str):
+    """Show details of a specific toolgroup at distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    toolgroups_get_response = client.tools.list()
+    # filter response to only include provided toolgroup_id
+    toolgroups_get_response = [
+        toolgroup for toolgroup in toolgroups_get_response if toolgroup.toolgroup_id == toolgroup_id
+    ]
+    if len(toolgroups_get_response) == 0:
+        console.print(
+            f"Toolgroup {toolgroup_id} is not found at distribution endpoint. "
+            "Please ensure endpoint is serving specified toolgroup.",
+            style="bold red",
+        )
+        return
+
+    headers = sorted(toolgroups_get_response[0].__dict__.keys())
+    table = Table()
+    for header in headers:
+        table.add_column(header)
+
+    for toolgroup in toolgroups_get_response:
+        row = [str(getattr(toolgroup, header)) for header in headers]
+        table.add_row(*row)
+    console.print(table)
+
+
+@click.command(name="register", help="Register a new toolgroup at distribution endpoint")
+@click.help_option("-h", "--help")
+@click.argument("toolgroup_id")
+@click.option("--provider-id", help="Provider ID for the toolgroup", default=None)
+@click.option("--mcp-endpoint", help="JSON mcp_config for the toolgroup", default=None)
+@click.option("--args", help="JSON args for the toolgroup", default=None)
+@click.pass_context
+@handle_client_errors("register toolgroup")
+def register_toolgroup(
+    ctx,
+    toolgroup_id: str,
+    provider_id: str | None,
+    mcp_endpoint: str | None,
+    args: str | None,
+):
+    """Register a new toolgroup at distribution endpoint"""
+    client = ctx.obj["client"]
+    console = Console()
+
+    _mcp_endpoint: toolgroup_register_params.McpEndpoint | NotGiven = NOT_GIVEN
+    if mcp_endpoint:
+        _mcp_endpoint = toolgroup_register_params.McpEndpoint(uri=mcp_endpoint)
+
+    response = client.toolgroups.register(
+        toolgroup_id=toolgroup_id,
+        provider_id=provider_id,
+        args=args,
+        mcp_endpoint=_mcp_endpoint,
+    )
+    if response:
+        console.print(f"[green]Successfully registered toolgroup {toolgroup_id}[/green]")
+
+
+@click.command(name="unregister", help="Unregister a toolgroup from distribution endpoint")
+@click.help_option("-h", "--help")
+@click.argument("toolgroup_id")
+@click.pass_context
+@handle_client_errors("unregister toolgroup")
+def unregister_toolgroup(ctx, toolgroup_id: str):
+    client = ctx.obj["client"]
+    console = Console()
+
+    response = client.toolgroups.unregister(toolgroup_id=toolgroup_id)
+    if response:
+        console.print(f"[green]Successfully deleted toolgroup {toolgroup_id}[/green]")
+
+
+# Register subcommands
+toolgroups.add_command(list_toolgroups)
+toolgroups.add_command(get_toolgroup)
+toolgroups.add_command(register_toolgroup)
+toolgroups.add_command(unregister_toolgroup)

ogx_open_client/lib/cli/vector_stores/__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 .vector_stores import vector_stores
+
+__all__ = ["vector_stores"]

ogx_open_client/lib/cli/vector_stores/vector_stores.py

@@ -0,0 +1,138 @@
+# 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 rich.table import Table
+
+from ..common.utils import handle_client_errors
+
+
+@click.group()
+@click.help_option("-h", "--help")
+def vector_stores():
+    """Manage vector stores."""
+
+
+@click.command("list")
+@click.help_option("-h", "--help")
+@click.pass_context
+@handle_client_errors("list vector stores")
+def list(ctx):
+    """Show available vector stores on distribution endpoint"""
+
+    client = ctx.obj["client"]
+    console = Console()
+    vector_stores_list_response = client.vector_stores.list()
+
+    if vector_stores_list_response:
+        table = Table()
+        # Add our specific columns
+        table.add_column("id")
+        table.add_column("name")
+        table.add_column("provider_id")
+        table.add_column("provider_vector_store_id")
+        table.add_column("params")
+
+        for item in vector_stores_list_response:
+            metadata = dict(item.metadata or {})
+            provider_id = str(metadata.pop("provider_id", ""))
+            provider_vector_store_id = str(
+                metadata.pop("provider_vector_store_id", metadata.pop("provider_vector_db_id", ""))
+            )
+            params = {
+                "status": item.status,
+                "usage_bytes": item.usage_bytes,
+                "embedding_model": metadata.pop("embedding_model", None),
+                "embedding_dimension": metadata.pop("embedding_dimension", None),
+                "metadata": metadata or None,
+            }
+            params = yaml.dump(params, default_flow_style=False)
+
+            table.add_row(
+                str(item.id),
+                str(item.name or ""),
+                provider_id,
+                provider_vector_store_id,
+                params,
+            )
+
+        console.print(table)
+
+
+@vector_stores.command()
+@click.help_option("-h", "--help")
+@click.argument("vector-db-id")
+@click.option("--provider-id", help="Provider ID for the vector db", default=None)
+@click.option("--provider-vector-db-id", help="Provider's vector db ID", default=None)
+@click.option(
+    "--embedding-model",
+    type=str,
+    help="Embedding model (for vector type)",
+    default="all-MiniLM-L6-v2",
+)
+@click.option(
+    "--embedding-dimension",
+    type=int,
+    help="Embedding dimension (for vector type)",
+    default=384,
+)
+@click.pass_context
+@handle_client_errors("register vector store")
+def register(
+    ctx,
+    vector_db_id: str,
+    provider_id: str | None,
+    provider_vector_db_id: str | None,
+    embedding_model: str | None,
+    embedding_dimension: int | None,
+):
+    """Create a new vector store"""
+    client = ctx.obj["client"]
+
+    metadata = {"vector_db_id": vector_db_id}
+    if provider_id:
+        metadata["provider_id"] = provider_id
+    if provider_vector_db_id:
+        metadata["provider_vector_store_id"] = provider_vector_db_id
+    if embedding_model:
+        metadata["embedding_model"] = embedding_model
+    if embedding_dimension is not None:
+        metadata["embedding_dimension"] = embedding_dimension
+
+    response = client.vector_stores.create(name=vector_db_id, metadata=metadata)
+    if response:
+        click.echo(yaml.dump(response.to_dict()))
+
+
+@vector_stores.command()
+@click.help_option("-h", "--help")
+@click.argument("vector-db-id")
+@click.pass_context
+@handle_client_errors("delete vector store")
+def unregister(ctx, vector_db_id: str):
+    """Delete a vector store"""
+    client = ctx.obj["client"]
+    vector_store_id = vector_db_id
+    for store in client.vector_stores.list():
+        if store.id == vector_db_id or store.name == vector_db_id:
+            vector_store_id = store.id
+            break
+        metadata = store.metadata or {}
+        if metadata.get("vector_db_id") == vector_db_id:
+            vector_store_id = store.id
+            break
+
+    client.vector_stores.delete(vector_store_id=vector_store_id)
+    click.echo(f"Vector store '{vector_db_id}' deleted successfully")
+
+
+# Register subcommands
+vector_stores.add_command(list)
+vector_stores.add_command(register)
+vector_stores.add_command(unregister)

ogx_open_client/lib/inference/__init__.py

@@ -0,0 +1,5 @@
+# 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.

ogx_open_client/lib/inference/event_logger.py

@@ -0,0 +1,58 @@
+# 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 collections.abc import Generator
+
+from termcolor import cprint
+
+from ...models.open_ai_chat_completion_chunk import OpenAIChatCompletionChunk as ChatCompletionChunk
+
+
+class InferenceStreamPrintableEvent:
+    def __init__(
+        self,
+        content: str = "",
+        end: str = "\n",
+        color="white",
+    ):
+        self.content = content
+        self.color = color
+        self.end = "\n" if end is None else end
+
+    def print(self, flush=True):
+        cprint(f"{self.content}", color=self.color, end=self.end, flush=flush)
+
+
+class InferenceStreamLogEventPrinter:
+    def __init__(self):
+        self.is_thinking = False
+
+    def yield_printable_events(
+        self, chunk: ChatCompletionChunk
+    ) -> Generator[InferenceStreamPrintableEvent, None, None]:
+        if hasattr(chunk, "choices") and len(chunk.choices) > 0:
+            yield from self._handle_chat_completion_chunk(chunk)
+
+    def _handle_chat_completion_chunk(
+        self, chunk: ChatCompletionChunk
+    ) -> Generator[InferenceStreamPrintableEvent, None, None]:
+        choice = chunk.choices[0]
+        delta = choice.delta
+        if delta:
+            if delta.role:
+                yield InferenceStreamPrintableEvent(f"{delta.role}> ", color="cyan", end="")
+            if delta.content:
+                yield InferenceStreamPrintableEvent(delta.content, color="yellow", end="")
+            if choice.finish_reason:
+                if choice.finish_reason == "length":
+                    yield InferenceStreamPrintableEvent("<truncated>", color="red", end="")
+                yield InferenceStreamPrintableEvent()
+
+
+class EventLogger:
+    def log(self, event_generator):
+        printer = InferenceStreamLogEventPrinter()
+        for chunk in event_generator:
+            yield from printer.yield_printable_events(chunk)

ogx_open_client/lib/inference/utils.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 base64
+import pathlib
+
+
+class MessageAttachment:
+    # https://developer.mozilla.org/en-US/docs/Glossary/Base64
+    @classmethod
+    def base64(cls, file_path: str) -> str:
+        path = pathlib.Path(file_path)
+        return base64.b64encode(path.read_bytes()).decode("utf-8")
+
+    # https://developer.mozilla.org/en-US/docs/Web/URI/Schemes/data
+    @classmethod
+    def data_url(cls, media_type: str, file_path: str) -> str:
+        return f"data:{media_type};base64,{cls.base64(file_path)}"

ogx_open_client/lib/stream_printer.py

@@ -0,0 +1,30 @@
+# 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 .agents.event_logger import TurnStreamEventPrinter
+from .inference.event_logger import InferenceStreamLogEventPrinter
+
+
+class EventStreamPrinter:
+    @classmethod
+    def gen(cls, event_generator):
+        inference_printer = None
+        turn_printer = None
+        for chunk in event_generator:
+            if not hasattr(chunk, "event"):
+                raise ValueError(f"Unexpected chunk without event: {chunk}")
+
+            event = chunk.event
+            if hasattr(event, "event_type"):
+                if not inference_printer:
+                    inference_printer = InferenceStreamLogEventPrinter()
+                yield from inference_printer.yield_printable_events(chunk)
+            elif hasattr(event, "payload") and hasattr(event.payload, "event_type"):
+                if not turn_printer:
+                    turn_printer = TurnStreamEventPrinter()
+                yield from turn_printer.yield_printable_events(chunk)
+            else:
+                raise ValueError(f"Unsupported event: {event}")

ogx_open_client/lib/tools/__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 .mcp_oauth import get_oauth_token_for_mcp_server
+
+__all__ = ["get_oauth_token_for_mcp_server"]