ibm-watsonx-orchestrate 1.0.0__py3-none-any.whl

ibm_watsonx_orchestrate/agent_builder/utils/pydantic_utils.py

@@ -0,0 +1,149 @@
+import json
+from pathlib import Path
+from typing import Any, Dict, Type, Literal
+
+import jsonschema
+from pydantic import BaseModel
+from pydantic.json_schema import DEFAULT_REF_TEMPLATE, GenerateJsonSchema, JsonSchemaMode
+from pydantic.main import IncEx
+from typing_extensions import Self
+
+from ibm_watsonx_orchestrate.agent_builder.tools import ToolRequestBody, ToolResponseBody, JsonSchemaObject
+
+
+def generate_schema_only_base_model(schema: ToolRequestBody | ToolResponseBody | JsonSchemaObject) -> Type[BaseModel]:
+    class SchemaOnlyBaseModel(BaseModel):
+        __primitive__: Any
+        model_config = {
+            'extra': 'allow'
+        }
+        """
+        The purpose of a SchemaOnlyBaseModel is to pass along the json schema represented by schema
+        to a langchain tool's arg_schema
+        :arg schema The json schema to emulate
+        :returns a fake BaseModel that only returns a json schema
+        """
+
+        def __init__(self, /, *args, **kwargs: Any) -> None:
+            if schema.type == 'object':
+                super().__init__(**kwargs)
+                for name, value in kwargs.items():
+                    setattr(self, name, value)
+            else:
+                kwargs={}
+                kwargs['__primitive__'] = args[0]
+                super().__init__(**kwargs)
+                setattr(self, '__primitive__', args[0])
+
+        @classmethod
+        def model_json_schema(cls, by_alias: bool = True, ref_template: str = DEFAULT_REF_TEMPLATE,
+                              schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema,
+                              mode: JsonSchemaMode = 'validation') -> dict[str, Any]:
+            return schema.model_dump(exclude_unset=True, exclude_none=True)
+
+        @classmethod
+        def schema(cls, by_alias: bool = True, ref_template: str = DEFAULT_REF_TEMPLATE) -> Dict[str, Any]:
+            return schema.model_dump(exclude_unset=True, exclude_none=True)
+
+        @classmethod
+        def schema_json(cls, *, by_alias: bool = True, ref_template: str = DEFAULT_REF_TEMPLATE,
+                        **dumps_kwargs: Any) -> str:
+            return schema.model_dump_json(exclude_unset=True, exclude_none=True)
+
+        @classmethod
+        def model_validate_json(cls, json_data: str | bytes | bytearray, *, strict: bool | None = None,
+                                context: Any | None = None) -> Self:
+            obj = json.loads(json_data)
+            jsonschema.validate(obj, schema=schema.model_dump(exclude_unset=True))
+            if schema.type == 'object':
+                return SchemaOnlyBaseModel(**obj)
+            else:
+                return SchemaOnlyBaseModel(obj)
+
+        @classmethod
+        def model_validate(cls, obj: Any, *, strict: bool | None = None, from_attributes: bool | None = None,
+                           context: Any | None = None) -> Self:
+            jsonschema.validate(obj, schema=schema.model_dump(exclude_unset=True))
+            if schema.type == 'object':
+                return SchemaOnlyBaseModel(**obj)
+            else:
+                return SchemaOnlyBaseModel(obj)
+
+        @classmethod
+        def model_validate_strings(cls, obj: Any, *, strict: bool | None = None, context: Any | None = None) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def validate(cls, value: Any) -> Self:
+            jsonschema.validate(value, schema=schema.model_dump(exclude_unset=True))
+            return SchemaOnlyBaseModel(**value)
+
+        @classmethod
+        def model_construct(cls, _fields_set: set[str] | None = None, **values: Any) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def model_parametrized_name(cls, params: tuple[type[Any], ...]) -> str:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def model_rebuild(cls, *, force: bool = False, raise_errors: bool = True, _parent_namespace_depth: int = 2,
+                          _types_namespace: Any = None) -> bool | None:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def parse_obj(cls, obj: Any) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def parse_raw(cls, b: str | bytes, *, content_type: str | None = None, encoding: str = 'utf8',
+                      proto: Any | None = None, allow_pickle: bool = False) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def parse_file(cls, path: str | Path, *, content_type: str | None = None, encoding: str = 'utf8',
+                       proto: Any | None = None, allow_pickle: bool = False) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def from_orm(cls, obj: Any) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def construct(cls, _fields_set: set[str] | None = None, **values: Any) -> Self:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        @classmethod
+        def update_forward_refs(cls, **localns: Any) -> None:
+            raise NotImplementedError('unimplemented for SchemaOnlyBaseModel')
+
+        def model_dump(self, *, mode: Literal['json', 'python'] | str = 'python', include: IncEx | None = None,
+                       exclude: IncEx | None = None, context: Any | None = None, by_alias: bool = False,
+                       exclude_unset: bool = False, exclude_defaults: bool = False, exclude_none: bool = False,
+                       round_trip: bool = False, warnings: bool | Literal['none', 'warn', 'error'] = True,
+                       serialize_as_any: bool = False) -> dict[str, Any]:
+            primitive = getattr(self, '__primitive__', None)
+            if primitive is not None:
+                return primitive
+            else:
+                return super().model_dump(mode=mode, include=include, exclude=exclude, context=context, by_alias=by_alias,
+                                          exclude_unset=exclude_unset, exclude_defaults=exclude_defaults,
+                                          exclude_none=exclude_none, round_trip=round_trip, warnings=warnings,
+                                          serialize_as_any=serialize_as_any)
+
+        def model_dump_json(self, *, indent: int | None = None, include: IncEx | None = None,
+                            exclude: IncEx | None = None, context: Any | None = None, by_alias: bool = False,
+                            exclude_unset: bool = False, exclude_defaults: bool = False, exclude_none: bool = False,
+                            round_trip: bool = False, warnings: bool | Literal['none', 'warn', 'error'] = True,
+                            serialize_as_any: bool = False) -> str:
+            primitive = getattr(self, '__primitive__')
+            if primitive is not None:
+                return json.dumps(primitive)
+            else:
+                return super().model_dump_json(indent=indent, include=include, exclude=exclude, context=context,
+                                           by_alias=by_alias, exclude_unset=exclude_unset,
+                                           exclude_defaults=exclude_defaults, exclude_none=exclude_none,
+                                           round_trip=round_trip, warnings=warnings, serialize_as_any=serialize_as_any)
+
+    return SchemaOnlyBaseModel
+

ibm_watsonx_orchestrate/cli/commands/agents/agents_command.py

@@ -0,0 +1,192 @@
+import typer
+from typing_extensions import Annotated, List
+from ibm_watsonx_orchestrate.cli.commands.agents.agents_controller import AgentsController
+from ibm_watsonx_orchestrate.agent_builder.agents.types import DEFAULT_LLM, AgentKind, AgentStyle, ExternalAgentAuthScheme, AgentProvider
+import json
+
+agents_app = typer.Typer(no_args_is_help=True)
+
+
+@agents_app.command(name="import", help='Import an agent definition into the active env from a file')
+def agent_import(
+    file: Annotated[
+        str,
+        typer.Option("--file", "-f", help="YAML file with agent definition"),
+    ],
+    app_id: Annotated[
+        str, typer.Option(
+            '--app-id', '-a',
+            help='The app id of the connection to associate with this external agent. An application connection represents the server authentication credentials needed to connection to this agent (for example Api Keys, Basic, Bearer or OAuth credentials).'
+        )
+    ] = None,
+):
+    agents_controller = AgentsController()
+    agent_specs = agents_controller.import_agent(file=file, app_id=app_id)
+    agents_controller.publish_or_update_agents(agent_specs)
+
+
+@agents_app.command(name="create", help='Create and import an agent into the active env')
+def agent_create(
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="Name of the agent you wish to create"),
+    ],
+    title: Annotated[
+        str,
+        typer.Option("--title", "-t", help="Title of the agent you wish to create. Only needed for External and Assistant Agents"),
+    ] = None,
+    kind: Annotated[
+        AgentKind,
+        typer.Option("--kind", "-k", help="The kind of agent you wish to create"),
+    ] = AgentKind.NATIVE,
+    api_url: Annotated[
+        str,
+        typer.Option("--api", "-a", help="External Api url your Agent will use"),
+    ] = None,
+    auth_scheme: Annotated[
+        ExternalAgentAuthScheme,
+        typer.Option("--auth-scheme", help="External Api auth schema to be used"),
+    ] = ExternalAgentAuthScheme.NONE,
+    provider: Annotated[
+        AgentProvider,
+        typer.Option("--provider", "-p", help="Agent Provider to be used.")
+    ] = AgentProvider.EXT_CHAT,
+    auth_config: Annotated[
+        str,
+        typer.Option(
+            "--auth-config",
+            help="Auth configuration to be used in JSON format (e.g., '{\"token\": \"test-api-key1\"')",
+        ),
+    ] = {},
+    tags: Annotated[
+        List[str],
+        typer.Option(
+            "--tags",
+            help="A list of tags for the agent. Format: --tags tag1 --tags tag2 ... Only needed for External and Assistant Agents",
+        ),
+    ] = None,
+    chat_params: Annotated[
+        str,
+        typer.Option(
+            "--chat-params",
+            help="Chat parameters in JSON format (e.g., '{\"stream\": true}'). Only needed for External and Assistant Agents",
+        ),
+    ] = None,
+    config: Annotated[
+        str,
+        typer.Option(
+            "--config",
+            help="Agent configuration in JSON format (e.g., '{\"hidden\": false, \"enable_cot\": false}')",
+        ),
+    ] = None,
+    nickname: Annotated[
+        str,
+        typer.Option("--nickname", help="Agent's nickname"),
+    ] = None,
+    app_id: Annotated[
+        str,
+        typer.Option("--app-id", help="Application ID for the agent"),
+    ] = None,
+    description: Annotated[
+        str,
+        typer.Option(
+            "--description",
+            help="Description of the agent",
+        ),
+    ] = None,
+    llm: Annotated[
+        str,
+        typer.Option(
+            "--llm",
+            help="The LLM used by the agent",
+        ),
+    ] = DEFAULT_LLM,
+    style: Annotated[
+        AgentStyle,
+        typer.Option("--style", help="The style of agent you wish to create"),
+    ] = AgentStyle.DEFAULT,
+    collaborators: Annotated[
+        List[str],
+        typer.Option(
+            "--collaborators",
+            help="A list of agent names you wish for the agent to be able to collaborate with. Format --colaborators agent1 --collaborators agent2 ...",
+        ),
+    ] = None,
+    tools: Annotated[
+        List[str],
+        typer.Option(
+            "--tools",
+            help="A list of tool names you wish for the agent to be able to utilise. Format --tools tool1 --tools agent2 ...",
+        ),
+    ] = None,
+    knowledge_base: Annotated[
+        List[str],
+        typer.Option(
+            "--knowledge-bases",
+            help="A list of knowlege bases names you wish for the agent to be able to utilise. Format --knowledge-bases base1 --knowledge-bases base2 ...",
+        ),
+    ] = None,
+    output_file: Annotated[
+        str,
+        typer.Option(
+            "--output",
+            "-o",
+            help="Write the agent definition out to a YAML (.yaml/.yml) file or a JSON (.json) file.",
+        ),
+    ] = None,
+):
+    chat_params_dict = json.loads(chat_params) if chat_params else {}
+    config_dict = json.loads(config) if config else {}
+    auth_config_dict = json.loads(auth_config) if auth_config else {}
+
+    agents_controller = AgentsController()
+    agent = agents_controller.generate_agent_spec(
+        name=name,
+        kind=kind,
+        description=description,
+        title=title,
+        api_url=api_url,
+        auth_scheme=auth_scheme,
+        auth_config=auth_config_dict,
+        provider=provider,
+        llm=llm,
+        style=style,
+        collaborators=collaborators,
+        tools=tools,
+        knowledge_base=knowledge_base,
+        tags=tags,
+        chat_params=chat_params_dict,
+        config=config_dict,
+        nickname=nickname,
+        app_id=app_id,
+        output_file=output_file,
+    )
+    agents_controller.publish_or_update_agents([agent])
+
+@agents_app.command(name="list", help='List all agents in the active env')
+def list_agents(
+    kind: Annotated[
+        AgentKind,
+        typer.Option("--kind", "-k", help="The kind of agent you wish to create"),
+    ] = None,
+    verbose: Annotated[
+        bool,
+        typer.Option("--verbose", "-v", help="List full details of all agents in json format"),
+    ] = False,
+):  
+    agents_controller = AgentsController()
+    agents_controller.list_agents(kind=kind, verbose=verbose)
+
+@agents_app.command(name="remove", help='Remove an agent from the active env')
+def remove_agent(
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="Name of the agent you wish to remove"),
+    ],
+    kind: Annotated[
+        AgentKind,
+        typer.Option("--kind", "-k", help="The kind of agent you wish to remove"),
+    ]
+):  
+    agents_controller = AgentsController()
+    agents_controller.remove_agent(name=name, kind=kind)