ibm-watsonx-orchestrate 1.4.2__py3-none-any.whl → 1.5.0b1__py3-none-any.whl

ibm_watsonx_orchestrate/__init__.py

@@ -5,7 +5,7 @@
 
 pkg_name = "ibm-watsonx-orchestrate"
 
-__version__ = "1.4.2"
+__version__ = "1.5.0b1"
 
 
 

ibm_watsonx_orchestrate/agent_builder/agents/types.py

@@ -3,12 +3,14 @@ import yaml
 from enum import Enum
 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.tools import BaseTool, PythonTool
 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
 
+from ibm_watsonx_orchestrate.agent_builder.tools.types import JsonSchemaObject
+
 # TO-DO: this is just a placeholder. Will update this later to align with backend
 DEFAULT_LLM = "watsonx/meta-llama/llama-3-1-70b-instruct"
 
@@ -78,6 +80,8 @@ class AgentSpec(BaseAgentSpec):
     kind: AgentKind = AgentKind.NATIVE
     llm: str = DEFAULT_LLM
     style: AgentStyle = AgentStyle.DEFAULT
+    custom_join_tool: str | PythonTool | None = None
+    structured_output: Optional[JsonSchemaObject] = None
     instructions: Annotated[Optional[str], Field(json_schema_extra={"min_length_str":1})] = None
     collaborators: Optional[List[str]] | Optional[List['BaseAgentSpec']] = []
     tools: Optional[List[str]] | Optional[List['BaseTool']] = []
@@ -117,6 +121,11 @@ def validate_agent_fields(values: dict) -> dict:
         if collaborator == name:
             raise ValueError(f"Circular reference detected. The agent '{name}' cannot contain itself as a collaborator")
 
+    if values.get("style") == AgentStyle.PLANNER:
+        if not values.get("custom_join_tool") and not values.get("structured_output"):
+            raise ValueError("Either 'custom_join_tool' or 'structured_output' must be provided for planner style agents.")
+        if values.get("custom_join_tool") and values.get("structured_output"):
+            raise ValueError("Only one of 'custom_join_tool' or 'structured_output' can be provided for planner style agents.")
 
     return values
 

ibm_watsonx_orchestrate/agent_builder/knowledge_bases/types.py

@@ -184,11 +184,24 @@ class ElasticSearchConnection(BaseModel):
     result_filter: Optional[list] = None 
     field_mapping: Optional[FieldMapping] = None
 
+class CustomSearchConnection(BaseModel):
+    """
+    example:
+    {
+        "url": "https://customsearch.xxxx.us-east.codeengine.appdomain.cloud",
+        "filter": "...",
+        "metadata": {...}
+    }
+    """
+    url: str
+    filter: Optional[str] = None
+    metadata: Optional[dict] = None
 
 class IndexConnection(BaseModel):
     connection_id: Optional[str] = None
     milvus: Optional[MilvusConnection] = None
     elastic_search: Optional[ElasticSearchConnection] = None
+    custom_search: Optional[CustomSearchConnection] = None
 
     
 class ConversationalSearchConfig(BaseModel):
@@ -201,6 +214,7 @@ class ConversationalSearchConfig(BaseModel):
     confidence_thresholds: Optional[ConfidenceThresholds] = None
 
 class KnowledgeBaseBuiltInVectorIndexConfig(BaseModel):
+    embeddings_model_name: Optional[str] = None
     chunk_size: Optional[int] = None
     chunk_overlap: Optional[int] = None
     limit: Optional[int] = None

ibm_watsonx_orchestrate/agent_builder/model_policies/__init__.py

@@ -0,0 +1 @@
+from .types import ModelPolicy, ModelPolicyInner, ModelPolicyRetry, ModelPolicyStrategy, ModelPolicyStrategyMode, ModelPolicyTarget

ibm_watsonx_orchestrate/{client → agent_builder}/model_policies/types.py

@@ -1,7 +1,7 @@
 from enum import Enum
-from typing import Optional, List, Dict, Any, Union
+from typing import List, Union
 
-from pydantic import Field, BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict
 
 class ModelPolicyStrategyMode(str, Enum):
     LOAD_BALANCED = "loadbalance"
@@ -17,14 +17,13 @@ class ModelPolicyRetry(BaseModel):
     on_status_codes: List[int] = None
 
 class ModelPolicyTarget(BaseModel):
-    model_id: str = None
-    weight: int = None
-
+    weight: float = None
+    model_name: str = None
 
 class ModelPolicyInner(BaseModel):
     strategy: ModelPolicyStrategy = None
     retry: ModelPolicyRetry = None
-    targets: List[Union[ModelPolicyTarget, 'ModelPolicyInner']] = None
+    targets: List[Union['ModelPolicyInner', ModelPolicyTarget]] = None
 
 
 class ModelPolicy(BaseModel):
@@ -32,5 +31,5 @@ class ModelPolicy(BaseModel):
 
     name: str
     display_name: str
-    policy: ModelPolicyInner
-
+    description: str
+    policy: ModelPolicyInner

ibm_watsonx_orchestrate/agent_builder/models/__init__.py

@@ -0,0 +1 @@
+from .types import VirtualModel

ibm_watsonx_orchestrate/{client → agent_builder}/models/types.py

@@ -1,6 +1,6 @@
 from typing import Optional, List, Dict, Any, Union
 from enum import Enum
-from pydantic import Field, BaseModel, ConfigDict
+from pydantic import Field, BaseModel, ConfigDict, model_validator
 
 class  ModelProvider(str, Enum):
     OPENAI = 'openai'
@@ -30,13 +30,40 @@ class  ModelProvider(str, Enum):
     def __repr__(self):
         return self.value
     
+    @classmethod
+    def has_value(cls, value):
+        return value in cls._value2member_map_ 
+
+class  ModelType(str, Enum):
+    CHAT = 'chat'
+    CHAT_VISION = 'chat_vision'
+    COMPLETION = 'completion'
+    EMBEDDING = 'embedding'
+
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return self.value
+    
+class  ModelType(str, Enum):
+    CHAT = 'chat'
+    CHAT_VISION = 'chat_vision'
+    COMPLETION = 'completion'
+    EMBEDDING = 'embedding'
+
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return self.value
 
 class ProviderConfig(BaseModel):
     # Required fields
     provider: Optional[str]=''
 
 
-    api_key: Optional[str] = '' # this is not optional
+    api_key: Optional[str] = None
     url_to_fetch: Optional[str] = Field(None, alias="urlToFetch")
 
     # Misc
@@ -123,7 +150,7 @@ class ProviderConfig(BaseModel):
     azure_extra_params: Optional[str] = Field(None, alias="azureExtraParams")
     azure_foundry_url: Optional[str] = Field(None, alias="azureFoundryUrl")
 
-    strict_open_ai_compliance: Optional[bool] = Field(False, alias="strictOpenAiCompliance")
+    strict_open_ai_compliance: Optional[bool] = Field(None, alias="strictOpenAiCompliance")
     mistral_fim_completion: Optional[str] = Field(None, alias="mistralFimCompletion")
 
     # Anthropic
@@ -140,6 +167,10 @@ class ProviderConfig(BaseModel):
     watsonx_version: Optional[str] = Field(None, alias="watsonxVersion")
     watsonx_space_id: Optional[str] = Field(None, alias="watsonxSpaceId")
     watsonx_project_id: Optional[str] = Field(None, alias="watsonxProjectId")
+    watsonx_deployment_id: Optional[str] = Field(None, alias="watsonxDeploymentId")
+    watsonx_cpd_url:Optional[str] = Field(None, alias="watsonxCpdUrl")
+    watsonx_cpd_username:Optional[str] = Field(None, alias="watsonxCpdUsername")
+    watsonx_cpd_password:Optional[str] = Field(None, alias="watsonxCpdPassword")
 
     model_config = {
         "populate_by_name": True,  # Replaces allow_population_by_field_name
@@ -147,19 +178,36 @@ class ProviderConfig(BaseModel):
         "json_schema_extra": lambda schema: schema.get("properties", {}).pop("provider", None)
     }
 
+    def update(self, new_config: "ProviderConfig") -> "ProviderConfig":
+        old_config_dict = dict(self)
+        new_config_dict = dict(new_config)
 
+        new_config_dict = {k:v for k, v in new_config_dict.items() if v is not None}
+        old_config_dict.update(new_config_dict)
 
-class CreateVirtualModel(BaseModel):
+        return ProviderConfig.model_validate(old_config_dict)
+
+
+class VirtualModel(BaseModel):
     model_config = ConfigDict(extra='allow')
 
     name: str
     display_name: Optional[str]
     description: Optional[str]
     config: Optional[dict] = None
-    provider_config: ProviderConfig
-    tags: List[str]
-    model_type: str = 'chat'
-
+    provider_config: Optional[ProviderConfig] = None
+    tags: List[str] = []
+    model_type: str = ModelType.CHAT
+    connection_id: Optional[str] = None
+
+    @model_validator(mode="before")
+    def validate_fields(cls, values):
+        if not values.get("display_name"):
+            values["display_name"] = values.get("name")
+        if not values.get("description"):
+            values["description"] = values.get("name")
+        
+        return values
 
 class ListVirtualModel(BaseModel):
     model_config = ConfigDict(extra='allow')
@@ -174,4 +222,4 @@ class ListVirtualModel(BaseModel):
     tags: Optional[List[str]] = None
     model_type: Optional[str] = None
 
-ANTHROPIC_DEFAULT_MAX_TOKENS = 4096
+ANTHROPIC_DEFAULT_MAX_TOKENS = 4096

ibm_watsonx_orchestrate/agent_builder/tools/python_tool.py

@@ -2,7 +2,7 @@ import importlib
 import inspect
 import json
 import os
-from typing import Callable, List
+from typing import Any, Callable, Dict, List, get_type_hints
 import logging
 
 import docstring_parser
@@ -13,12 +13,18 @@ from pydantic import TypeAdapter, BaseModel
 from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
 from ibm_watsonx_orchestrate.agent_builder.connections import ExpectedCredentials
 from .base_tool import BaseTool
-from .types import ToolSpec, ToolPermission, ToolRequestBody, ToolResponseBody, JsonSchemaObject, ToolBinding, \
+from .types import PythonToolKind, ToolSpec, ToolPermission, ToolRequestBody, ToolResponseBody, JsonSchemaObject, ToolBinding, \
     PythonToolBinding
 
 _all_tools = []
 logger = logging.getLogger(__name__)
 
+JOIN_TOOL_PARAMS = {
+    'original_query': str,
+    'task_results': Dict[str, Any],
+    'messages': List[Dict[str, Any]],
+}
+
 class PythonTool(BaseTool):
     def __init__(self, fn, spec: ToolSpec, expected_credentials: List[ExpectedCredentials]=None):
         BaseTool.__init__(self, spec=spec)
@@ -92,6 +98,31 @@ def _validate_input_schema(input_schema: ToolRequestBody) -> None:
         if not props.get(prop).type:
             logger.warning(f"Missing type hint for tool property '{prop}' defaulting to 'str'. To remove this warning add a type hint to the property in the tools signature. See Python docs for guidance: https://docs.python.org/3/library/typing.html")
 
+def _validate_join_tool_func(fn: Callable, sig: inspect.Signature | None = None, name: str | None = None) -> None:
+    if sig is None:
+        sig = inspect.signature(fn)
+    if name is None:
+        name = fn.__name__
+    
+    params = sig.parameters
+    type_hints = get_type_hints(fn)
+    
+    # Validate parameter order
+    actual_param_names = list(params.keys())
+    expected_param_names = list(JOIN_TOOL_PARAMS.keys())
+    if actual_param_names[:len(expected_param_names)] != expected_param_names:
+        raise ValueError(
+            f"Join tool function '{name}' has incorrect parameter names or order. Expected: {expected_param_names}, got: {actual_param_names}"
+        )
+    
+    # Validate the type hints
+    for param, expected_type in JOIN_TOOL_PARAMS.items():
+        if param not in type_hints:
+            raise ValueError(f"Join tool function '{name}' is missing type for parameter '{param}'")
+        actual_type = type_hints[param]
+        if actual_type != expected_type:
+            raise ValueError(f"Join tool function '{name}' has incorrect type for parameter '{param}'. Expected {expected_type}, got {actual_type}")
+
 def tool(
     *args,
     name: str = None,
@@ -100,7 +131,8 @@ def tool(
     output_schema: ToolResponseBody = None,
     permission: ToolPermission = ToolPermission.READ_ONLY,
     expected_credentials: List[ExpectedCredentials] = None,
-    display_name: str = None
+    display_name: str = None,
+    kind: PythonToolKind = PythonToolKind.TOOL,
 ) -> Callable[[{__name__, __doc__}], PythonTool]:
     """
     Decorator to convert a python function into a callable tool.
@@ -152,6 +184,11 @@ def tool(
         spec.binding.python.function = function_binding
 
         sig = inspect.signature(fn)
+        
+        # If the function is a join tool, validate its signature matches the expected parameters. If not, raise error with details.
+        if kind == PythonToolKind.JOIN_TOOL:
+            _validate_join_tool_func(fn, sig, spec.name)
+        
         if not input_schema:
             try:
                 input_schema_model: type[BaseModel] = create_schema_from_function(spec.name, fn, parse_docstring=True)
@@ -190,6 +227,12 @@ def tool(
 
         else:
             spec.output_schema = ToolResponseBody()
+        
+        # Validate the generated schema still conforms to the requirement for a join tool
+        if kind == PythonToolKind.JOIN_TOOL:
+            if not spec.is_custom_join_tool():
+                raise ValueError(f"Join tool '{spec.name}' does not conform to the expected join tool schema. Please ensure the input schema has the required fields: {JOIN_TOOL_PARAMS.keys()} and the output schema is a string.")
+            
         _all_tools.append(t)
         return t
 

ibm_watsonx_orchestrate/agent_builder/tools/types.py

@@ -10,6 +10,9 @@ class ToolPermission(str, Enum):
     READ_WRITE = 'read_write'
     ADMIN = 'admin'
 
+class PythonToolKind(str, Enum):
+    JOIN_TOOL = 'join_tool'
+    TOOL = 'tool'
 
 class JsonSchemaObject(BaseModel):
     model_config = ConfigDict(extra='allow')
@@ -163,7 +166,6 @@ class ToolBinding(BaseModel):
             raise ValueError("Only one binding can be set")
         return self
 
-
 class ToolSpec(BaseModel):
     name: str
     display_name: str | None = None
@@ -174,3 +176,47 @@ class ToolSpec(BaseModel):
     binding: ToolBinding = None
     toolkit_id: str | None = None
 
+    def is_custom_join_tool(self) -> bool:
+        if self.binding.python is None:
+            return False
+
+        # The code below validates the input schema to have the following structure:
+        # {
+        #     "type": "object",
+        #     "properties": {
+        #         "messages": {
+        #             "type": "array",
+        #             "items": {
+        #                 "type": "object",
+        #             },
+        #         },
+        #         "task_results": {
+        #             "type": "object",
+        #         },
+        #         "original_query": {
+        #             "type": "string",
+        #         },
+        #     },
+        #     "required": {"original_query", "task_results", "messages"},
+        # }
+
+        input_schema = self.input_schema
+        if input_schema.type != 'object':
+            return False
+
+        required_fields = {"original_query", "task_results", "messages"}
+        if input_schema.required is None or set(input_schema.required) != required_fields:
+            return False
+        if input_schema.properties is None or set(input_schema.properties.keys()) != required_fields:
+            return False
+
+        if input_schema.properties["messages"].type != "array":
+            return False
+        if not input_schema.properties["messages"].items or input_schema.properties["messages"].items.type != "object":
+            return False
+        if input_schema.properties["task_results"].type != "object":
+            return False
+        if input_schema.properties["original_query"].type != "string":
+            return False
+
+        return True

ibm_watsonx_orchestrate/cli/commands/agents/agents_command.py

@@ -105,6 +105,20 @@ def agent_create(
         AgentStyle,
         typer.Option("--style", help="The style of agent you wish to create"),
     ] = AgentStyle.DEFAULT,
+    custom_join_tool: Annotated[
+        str | None,
+        typer.Option(
+            "--custom-join-tool",
+            help='The name of the python tool to be used by the agent to format and generate the final output. Only needed for "planner" style agents.',
+        ),
+    ] = None,
+    structured_output: Annotated[
+        str | None,
+        typer.Option(
+            "--structured-output",
+            help='A JSON Schema object that defines the desired structure of the agent\'s final output. Only needed for "planner" style agents.',
+        ),
+    ] = None,
     collaborators: Annotated[
         List[str],
         typer.Option(
@@ -138,6 +152,7 @@ def agent_create(
     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 {}
+    structured_output_dict = json.loads(structured_output) if structured_output else None
 
     agents_controller = AgentsController()
     agent = agents_controller.generate_agent_spec(
@@ -151,6 +166,8 @@ def agent_create(
         provider=provider,
         llm=llm,
         style=style,
+        custom_join_tool=custom_join_tool,
+        structured_output=structured_output_dict,
         collaborators=collaborators,
         tools=tools,
         knowledge_base=knowledge_base,