ibm-watsonx-orchestrate 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

ibm_watsonx_orchestrate/__init__.py

@@ -5,7 +5,7 @@
 
 pkg_name = "ibm-watsonx-orchestrate"
 
-__version__ = "1.1.0"
+__version__ = "1.3.0"
 
 
 

ibm_watsonx_orchestrate/agent_builder/agents/types.py

@@ -5,6 +5,7 @@ from typing import List, Optional, Dict
 from pydantic import BaseModel, model_validator, ConfigDict
 from ibm_watsonx_orchestrate.agent_builder.tools import BaseTool
 from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import KnowledgeBaseSpec
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.knowledge_base import KnowledgeBase
 from pydantic import Field, AliasChoices
 from typing import Annotated
 
@@ -85,6 +86,8 @@ class AgentSpec(BaseAgentSpec):
     def __init__(self, *args, **kwargs):
         if "tools" in kwargs and kwargs["tools"]:
             kwargs["tools"] = [x.__tool_spec__.name if isinstance(x, BaseTool) else x for x in kwargs["tools"]]
+        if "knowledge_base" in kwargs and kwargs["knowledge_base"]:
+            kwargs["knowledge_base"] = [x.name if isinstance(x, KnowledgeBase) else x for x in kwargs["knowledge_base"]]
         if "collaborators" in kwargs and kwargs["collaborators"]:
             kwargs["collaborators"] = [x.name if isinstance(x, BaseAgentSpec) else x for x in kwargs["collaborators"]]
         super().__init__(*args, **kwargs)
@@ -101,7 +104,7 @@ class AgentSpec(BaseAgentSpec):
 
 def validate_agent_fields(values: dict) -> dict:
     # Check for empty strings or whitespace
-    for field in ["id", "name", "kind", "description", "collaborators", "tools"]:
+    for field in ["id", "name", "kind", "description", "collaborators", "tools", "knowledge_base"]:
         value = values.get(field)
         if value and not str(value).strip():
             raise ValueError(f"{field} cannot be empty or just whitespace")

ibm_watsonx_orchestrate/agent_builder/knowledge_bases/knowledge_base.py

@@ -1,7 +1,7 @@
 import json
 from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
-from .types import KnowledgeBaseSpec
-
+from .types import KnowledgeBaseSpec, KnowledgeBaseKind
+from pydantic import model_validator
 
 class KnowledgeBase(KnowledgeBaseSpec):
 
@@ -24,4 +24,17 @@ class KnowledgeBase(KnowledgeBaseSpec):
         return f"KnowledgeBase(id='{self.id}', name='{self.name}', description='{self.description}')"
 
     def __str__(self):
-        return self.__repr__()
+        return self.__repr__()
+    
+    # Not a model validator since we only want to validate this on import
+    def validate_documents_or_index_exists(self):
+        if self.documents and self.conversational_search_tool and self.conversational_search_tool.index_config or \
+            (not self.documents and (not self.conversational_search_tool or not self.conversational_search_tool.index_config)):
+            raise ValueError("Must provide either \"documents\" or \"conversational_search_tool.index_config\", but not both")
+        return self
+    
+    @model_validator(mode="after")
+    def validate_kind(self):
+        if self.kind != KnowledgeBaseKind.KNOWLEDGE_BASE:
+            raise ValueError(f"The specified kind '{self.kind}' cannot be used to create a knowledge base")
+        return self

ibm_watsonx_orchestrate/agent_builder/knowledge_bases/knowledge_base_requests.py

@@ -1,12 +1,12 @@
 import json
 from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
-from .types import CreateKnowledgeBase, PatchKnowledgeBase, KnowledgeBaseKind
+from .types import KnowledgeBaseSpec, PatchKnowledgeBase, KnowledgeBaseKind
 
 
-class KnowledgeBaseCreateRequest(CreateKnowledgeBase):
+class KnowledgeBaseCreateRequest(KnowledgeBaseSpec):
 
     @staticmethod
-    def from_spec(file: str) -> 'CreateKnowledgeBase':
+    def from_spec(file: str) -> 'KnowledgeBaseSpec':
         with open(file, 'r') as f:
             if file.endswith('.yaml') or file.endswith('.yml'):
                 content = yaml_safe_load(f)
@@ -15,20 +15,10 @@ class KnowledgeBaseCreateRequest(CreateKnowledgeBase):
             else:
                 raise ValueError('file must end in .json, .yaml, or .yml')
             
-            if (content.get('documents') and content.get("conversational_search_tool", {}).get("index_config")) or \
-                  (not content.get('documents') and not content.get("conversational_search_tool", {}).get("index_config")):
-                raise ValueError("Must provide either \"documents\" or \"conversational_search_tool.index_config\", but not both")
-            
             if not content.get("spec_version"):
                 raise ValueError(f"Field 'spec_version' not provided. Please ensure provided spec conforms to a valid spec format")
             
-            if not content.get("kind"):
-                raise ValueError(f"Field 'kind' not provided. Should be 'knowledge_base'")
-
-            if content.get("kind") != KnowledgeBaseKind.KNOWLEDGE_BASE:
-                raise ValueError(f"Field 'kind' should be 'knowledge_base', but is set to '{content.get('kind')}'")
-            
-            knowledge_base = CreateKnowledgeBase.model_validate(content)
+            knowledge_base = KnowledgeBaseSpec.model_validate(content)
 
         return knowledge_base
     
@@ -48,12 +38,6 @@ class KnowledgeBaseUpdateRequest(PatchKnowledgeBase):
             if not content.get("spec_version"):
                 raise ValueError(f"Field 'spec_version' not provided. Please ensure provided spec conforms to a valid spec format")
             
-            if not content.get("kind"):
-                raise ValueError(f"Field 'kind' not provided. Should be 'knowledge_base'")
-
-            if content.get("kind") != KnowledgeBaseKind.KNOWLEDGE_BASE:
-                raise ValueError(f"Field 'kind' should be 'knowledge_base', but is set to '{content.get('kind')}'")
-            
             patch = PatchKnowledgeBase.model_validate(content)
 
         return patch

ibm_watsonx_orchestrate/agent_builder/knowledge_bases/types.py

@@ -206,16 +206,6 @@ class KnowledgeBaseBuiltInVectorIndexConfig(BaseModel):
     chunk_overlap: Optional[int] = None
     limit: Optional[int] = None
     
-class CreateKnowledgeBase(BaseModel):
-    """request payload schema"""
-    name: Optional[str] = None
-    description: Optional[str] = None
-    documents: list[str] = None
-    vector_index: Optional[KnowledgeBaseBuiltInVectorIndexConfig] = None
-    conversational_search_tool: Optional[ConversationalSearchConfig] = None
-    prioritize_built_in_index: Optional[bool] = None
-    
-    
 class PatchKnowledgeBase(BaseModel):
     """request payload schema"""
     description: Optional[str] = None
@@ -224,20 +214,21 @@ class PatchKnowledgeBase(BaseModel):
     prioritize_built_in_index: Optional[bool] = None
     representation: Optional[KnowledgeBaseRepresentation] = None
     
-    
 class KnowledgeBaseSpec(BaseModel):
     """Schema for a complete knowledge-base."""
     spec_version: SpecVersion = None
     kind: KnowledgeBaseKind = KnowledgeBaseKind.KNOWLEDGE_BASE
     id: Optional[UUID] = None 
     tenant_id: Optional[str] = None
-    name: Optional[str] = None
+    name: str
     description: Optional[str] = None
     vector_index: Optional[KnowledgeBaseBuiltInVectorIndexConfig] = None
     conversational_search_tool: Optional[ConversationalSearchConfig] | Optional[UUID] = None
     prioritize_built_in_index: Optional[bool] = None
+    representation: Optional[KnowledgeBaseRepresentation] = None
     vector_index_id: Optional[UUID] = None 
     created_by: Optional[str] = None
     created_on: Optional[datetime] = None 
     updated_at: Optional[datetime] = None
-
+    # For import/update
+    documents: list[str] = None

ibm_watsonx_orchestrate/agent_builder/tools/openapi_tool.py

@@ -115,17 +115,9 @@ def create_openapi_json_tool(
         raise ValueError(
             f"Path {http_path} did not have an http_method {http_method}. Available methods are {list(route.keys())}")
 
-    operation_id = (
-        re.sub(
-            '_+',
-            '_',
-            re.sub(
-                r'[^a-zA-Z_]',
-                '_',
-                route_spec.get('operationId', None)
-            )
-        )
-    ) if route_spec.get('operationId', None) is not None else None
+    operation_id = re.sub( r'(\W|_)+', '_', route_spec.get('operationId') ) \
+                     if route_spec.get('operationId', None) else None
+
     spec_name = name or operation_id
     spec_permission = permission or _action_to_perm(route_spec.get('x-ibm-operation', {}).get('action'))
     if spec_name is None:
@@ -142,7 +134,7 @@ def create_openapi_json_tool(
         description=spec_description,
         permission=spec_permission
     )
-    
+
     spec.input_schema = input_schema or ToolRequestBody(
         type='object',
         properties={},

ibm_watsonx_orchestrate/agent_builder/tools/types.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, Any, Dict, Literal, Optional
+from typing import List, Any, Dict, Literal, Optional, Union
 
 from pydantic import BaseModel, model_validator, ConfigDict, Field, AliasChoices
 
@@ -14,7 +14,7 @@ class ToolPermission(str, Enum):
 class JsonSchemaObject(BaseModel):
     model_config = ConfigDict(extra='allow')
 
-    type: Optional[Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null']] = None
+    type: Optional[Union[Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null'], List[Literal['object', 'string', 'number', 'integer', 'boolean', 'array', 'null']]]] = None
     title: str | None = None
     description: str | None = None
     properties: Optional[Dict[str, 'JsonSchemaObject']] = None
@@ -34,13 +34,19 @@ class JsonSchemaObject(BaseModel):
     aliasName: str | None = None
     "Runtime feature where the sdk can provide the original name of a field before prefixing"
 
+    @model_validator(mode='after')
+    def normalize_type_field(self) -> 'JsonSchemaObject':
+        if isinstance(self.type, list):
+            self.type = self.type[0]
+        return self
+
 
 class ToolRequestBody(BaseModel):
     model_config = ConfigDict(extra='allow')
 
     type: Literal['object']
     properties: Dict[str, JsonSchemaObject]
-    required: List[str]
+    required: Optional[List[str]] = None
 
 
 class ToolResponseBody(BaseModel):
@@ -52,7 +58,7 @@ class ToolResponseBody(BaseModel):
     items: JsonSchemaObject = None
     uniqueItems: bool = None
     anyOf: List['JsonSchemaObject'] = None
-    required: List[str] = None
+    required: Optional[List[str]] = None
 
 class OpenApiSecurityScheme(BaseModel):
     type: Literal['apiKey', 'http', 'oauth2', 'openIdConnect']
@@ -128,6 +134,10 @@ class SkillToolBinding(BaseModel):
 class ClientSideToolBinding(BaseModel):
     pass
 
+class McpToolBinding(BaseModel):
+    server_url: Optional[str] = None
+    source: str
+    connections: Dict[str, str]
 
 class ToolBinding(BaseModel):
     openapi: OpenApiToolBinding = None
@@ -135,6 +145,7 @@ class ToolBinding(BaseModel):
     wxflows: WxFlowsToolBinding = None
     skill: SkillToolBinding = None
     client_side: ClientSideToolBinding = None
+    mcp: McpToolBinding = None
 
     @model_validator(mode='after')
     def validate_binding_type(self) -> 'ToolBinding':
@@ -143,7 +154,8 @@ class ToolBinding(BaseModel):
             self.python is not None,
             self.wxflows is not None,
             self.skill is not None,
-            self.client_side is not None
+            self.client_side is not None,
+            self.mcp is not None
         ]
         if sum(bindings) == 0:
             raise ValueError("One binding must be set")
@@ -159,4 +171,5 @@ class ToolSpec(BaseModel):
     input_schema: ToolRequestBody = None
     output_schema: ToolResponseBody = None
     binding: ToolBinding = None
+    toolkit_id: str | None = None
 

ibm_watsonx_orchestrate/cli/commands/agents/agents_controller.py

@@ -11,6 +11,7 @@ from copy import deepcopy
 
 from typing import Iterable, List
 from ibm_watsonx_orchestrate.cli.commands.tools.tools_controller import import_python_tool
+from ibm_watsonx_orchestrate.cli.commands.knowledge_bases.knowledge_bases_controller import import_python_knowledge_base
 
 from ibm_watsonx_orchestrate.agent_builder.agents import (
     Agent,
@@ -33,6 +34,7 @@ logger = logging.getLogger(__name__)
 def import_python_agent(file: str) -> List[Agent | ExternalAgent | AssistantAgent]:
     # Import tools
     import_python_tool(file)
+    import_python_knowledge_base(file)
 
     file_path = Path(file)
     file_directory = file_path.parent

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_command.py

@@ -9,7 +9,7 @@ knowledge_bases_app = typer.Typer(no_args_is_help=True)
 def knowledge_base_import(
     file: Annotated[
         str,
-        typer.Option("--file", "-f", help="YAML file with knowledge base definition"),
+        typer.Option("--file", "-f", help="YAML, JSON or Python file with knowledge base definition(s)"),
     ],
     app_id: Annotated[
         str, typer.Option(
@@ -25,7 +25,7 @@ def knowledge_base_import(
 def knowledge_base_patch(
     file: Annotated[
         str,
-        typer.Option("--file", "-f", help="YAML file with knowledge base definition"),
+        typer.Option("--file", "-f", help="YAML or JSON file with knowledge base definition"),
     ],
     name: Annotated[
         str,

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_controller.py

@@ -3,17 +3,44 @@ import json
 import rich
 import requests
 import logging
+import importlib
+import inspect
+from pathlib import Path
+from typing import List
 
-from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.knowledge_base_requests import KnowledgeBaseCreateRequest, KnowledgeBaseUpdateRequest
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.knowledge_base_requests import KnowledgeBaseUpdateRequest
 from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.knowledge_base import KnowledgeBase
 from ibm_watsonx_orchestrate.client.knowledge_bases.knowledge_base_client import KnowledgeBaseClient
 from ibm_watsonx_orchestrate.client.base_api_client import ClientAPIException
 from ibm_watsonx_orchestrate.client.connections import get_connections_client
-
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
 
 logger = logging.getLogger(__name__)
 
+def import_python_knowledge_base(file: str) -> List[KnowledgeBase]:
+    file_path = Path(file)
+    file_directory = file_path.parent
+    file_name = file_path.stem
+    sys.path.append(str(file_directory))
+    module = importlib.import_module(file_name)
+    del sys.path[-1]
+
+    knowledge_bases = []
+    for _, obj in inspect.getmembers(module):
+        if isinstance(obj, KnowledgeBase):
+            knowledge_bases.append(obj)
+    return knowledge_bases
+
+def parse_file(file: str) -> List[KnowledgeBase]:
+    if file.endswith('.yaml') or file.endswith('.yml') or file.endswith(".json"):
+        knowledge_base = KnowledgeBase.from_spec(file=file)
+        return [knowledge_base]
+    elif file.endswith('.py'):
+        knowledge_bases = import_python_knowledge_base(file)
+        return knowledge_bases
+    else:
+        raise ValueError("file must end in .json, .yaml, .yml or .py")
+
 def to_column_name(col: str):
     return " ".join([word.capitalize() if not word[0].isupper() else word for word in col.split("_")])
 
@@ -34,45 +61,43 @@ class KnowledgeBaseController:
     
     def import_knowledge_base(self, file: str, app_id: str):
         client = self.get_client()
-        create_request = KnowledgeBaseCreateRequest.from_spec(file=file)
 
-        try:
-            if create_request.documents:
-                try:
+        knowledge_bases = parse_file(file=file)
+        for kb in knowledge_bases:
+            try:
+                kb.validate_documents_or_index_exists()
+                if kb.documents:
                     file_dir = "/".join(file.split("/")[:-1])
-                    files = [('files', (get_file_name(file_path), open(file_path if file_path.startswith("/") else f"{file_dir}/{file_path}", 'rb'))) for file_path in create_request.documents]
-                except Exception as e:
-                    logger.error(f"Error importing knowledge base: {str(e).replace('[Errno 2] ', '')}")
-                    sys.exit(1);
-
-                payload = create_request.model_dump(exclude_none=True);
-                payload.pop('documents');
-
-                client.create_built_in(payload=payload, files=files)
-            else:
-                if len(create_request.conversational_search_tool.index_config) != 1:
-                    raise ValueError(f"Must provide exactly one conversational_search_tool.index_config. Provided {len(create_request.conversational_search_tool.index_config)}.")
+                    files = [('files', (get_file_name(file_path), open(file_path if file_path.startswith("/") else (file_path if not file_dir else f"{file_dir}/{file_path}"), 'rb'))) for file_path in kb.documents]
+                    
+                    payload = kb.model_dump(exclude_none=True);
+                    payload.pop('documents');
+
+                    client.create_built_in(payload=payload, files=files)
+                else:
+                    if len(kb.conversational_search_tool.index_config) != 1:
+                        raise ValueError(f"Must provide exactly one conversational_search_tool.index_config. Provided {len(kb.conversational_search_tool.index_config)}.")
+                    
+                    if app_id:
+                        connections_client = get_connections_client()
+                        connection_id = None
+                        if app_id is not None:
+                            connections = connections_client.get_draft_by_app_id(app_id=app_id)
+                            if not connections:
+                                logger.error(f"No connection exists with the app-id '{app_id}'")
+                                exit(1)
+
+                            connection_id = connections.connection_id
+                            kb.conversational_search_tool.index_config[0].connection_id = connection_id
+
+                    client.create(payload=kb.model_dump(exclude_none=True))
                 
-                if app_id:
-                    connections_client = get_connections_client()
-                    connection_id = None
-                    if app_id is not None:
-                        connections = connections_client.get_draft_by_app_id(app_id=app_id)
-                        if not connections:
-                            logger.error(f"No connection exists with the app-id '{app_id}'")
-                            exit(1)
-
-                        connection_id = connections.connection_id
-                        create_request.conversational_search_tool.index_config[0].connection_id = connection_id
-
-                client.create(payload=create_request.model_dump(exclude_none=True))
-            
-            logger.info(f"Successfully imported knowledge base '{create_request.name}'")
-        except ClientAPIException as e:
-            if "duplicate key value violates unique constraint" in e.response.text:
-                logger.error(f"A knowledge base with the name '{create_request.name}' already exists. Failed to import knowledge base")
-            else:
-                logger.error(f"Error importing knowledge base '{create_request.name}\n' {e.response.text}")
+                logger.info(f"Successfully imported knowledge base '{kb.name}'")
+            except ClientAPIException as e:
+                if "duplicate key value violates unique constraint" in e.response.text:
+                    logger.error(f"A knowledge base with the name '{kb.name}' already exists. Failed to import knowledge base")
+                else:
+                    logger.error(f"Error importing knowledge base '{kb.name}\n' {e.response.text}")
     
     def get_id(
         self, id: str, name: str

ibm_watsonx_orchestrate/cli/commands/toolkit/toolkit_command.py

@@ -0,0 +1,71 @@
+import typer
+from typing import List
+from typing_extensions import Annotated, Optional
+from ibm_watsonx_orchestrate.cli.commands.toolkit.toolkit_controller import ToolkitController, ToolkitKind
+
+toolkits_app = typer.Typer(no_args_is_help=True)
+
+@toolkits_app.command(name="import")
+def import_toolkit(
+    kind: Annotated[
+        ToolkitKind,
+        typer.Option("--kind", "-k", help="Kind of toolkit, currently only MCP is supported"),
+    ],
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="Name of the toolkit"),
+    ],
+    description: Annotated[
+        str,
+        typer.Option("--description", help="Description of the toolkit"),
+    ],
+    package_root: Annotated[
+        str,
+        typer.Option("--package-root", "-p", help="Root directory of the MCP server package"),
+    ],
+    command: Annotated[
+        str,
+        typer.Option(
+            "--command", 
+            help="Command to start the MCP server. Can be a string (e.g. 'node dist/index.js --transport stdio') "
+                "or a JSON-style list of arguments (e.g. '[\"node\", \"dist/index.js\", \"--transport\", \"stdio\"]'). "
+                "The first argument will be used as the executable, the rest as its arguments."
+        ),
+    ],
+    tools: Annotated[
+        Optional[str],
+        typer.Option("--tools", "-t", help="Comma-separated list of tools to import. Or you can use `*` to use all tools"),
+    ] = None,
+    app_id: Annotated[
+        List[str],
+        typer.Option(
+            "--app-id", "-a", 
+            help='The app id of the connection to associate with this tool. A application connection represents the server authentication credentials needed to connect to this tool. Only type key_value is currently supported for MCP.'
+        )
+    ] = None
+):
+    if tools == "*":
+        tool_list = ["*"] # Wildcard to use all tools
+    elif tools:
+        tool_list = [tool.strip() for tool in tools.split(",")]
+    else:
+        tool_list = None
+
+    toolkit_controller = ToolkitController(
+    kind=kind,
+    name=name,
+    description=description,
+    package_root=package_root,
+    command=command,
+)
+    toolkit_controller.import_toolkit(tools=tool_list, app_id=app_id)
+
+@toolkits_app.command(name="remove")
+def remove_toolkit(
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="Name of the toolkit you wish to remove"),
+    ],
+):  
+    toolkit_controller = ToolkitController()
+    toolkit_controller.remove_toolkit(name=name)