ibm-watsonx-orchestrate 1.5.0b0__py3-none-any.whl → 1.5.0b1__py3-none-any.whl

ibm_watsonx_orchestrate/__init__.py

@@ -5,8 +5,7 @@
 
 pkg_name = "ibm-watsonx-orchestrate"
 
-__version__ = "1.5.0b0"
-
+__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):

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'
@@ -29,6 +29,10 @@ 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'
@@ -42,13 +46,24 @@ class  ModelType(str, Enum):
     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
@@ -135,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
@@ -152,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
@@ -159,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]
+    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')
@@ -186,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,

ibm_watsonx_orchestrate/cli/commands/agents/agents_controller.py

@@ -11,9 +11,12 @@ import logging
 from pathlib import Path
 from copy import deepcopy
 
-from typing import Iterable, List
+from typing import Iterable, List, TypeVar
+from ibm_watsonx_orchestrate.agent_builder.agents.types import AgentStyle
+from ibm_watsonx_orchestrate.agent_builder.tools.types import ToolSpec
 from ibm_watsonx_orchestrate.cli.commands.tools.tools_controller import import_python_tool, ToolsController
 from ibm_watsonx_orchestrate.cli.commands.knowledge_bases.knowledge_bases_controller import import_python_knowledge_base
+from ibm_watsonx_orchestrate.cli.commands.models.models_controller import import_python_model
 
 from ibm_watsonx_orchestrate.agent_builder.agents import (
     Agent,
@@ -34,10 +37,14 @@ from ibm_watsonx_orchestrate.utils.utils import check_file_in_zip
 
 logger = logging.getLogger(__name__)
 
+# Helper generic type for any agent
+AnyAgentT = TypeVar("AnyAgentT", bound=Agent | ExternalAgent | AssistantAgent)
+
 def import_python_agent(file: str) -> List[Agent | ExternalAgent | AssistantAgent]:
     # Import tools
     import_python_tool(file)
     import_python_knowledge_base(file)
+    import_python_model(file)
 
     file_path = Path(file)
     file_directory = file_path.parent
@@ -90,6 +97,8 @@ def parse_create_native_args(name: str, kind: AgentKind, description: str | None
         "description": description,
         "llm": args.get("llm"),
         "style": args.get("style"),
+        "custom_join_tool": args.get("custom_join_tool"),
+        "structured_output": args.get("structured_output"),
     }
 
     collaborators = args.get("collaborators", [])
@@ -202,7 +211,7 @@ class AgentsController:
         return self.knowledge_base_client
     
     @staticmethod
-    def import_agent(file: str, app_id: str) -> Iterable:
+    def import_agent(file: str, app_id: str) -> List[Agent | ExternalAgent | AssistantAgent]:
         agents = parse_file(file)
         for agent in agents:
             if app_id and agent.kind != AgentKind.NATIVE and agent.kind != AgentKind.ASSISTANT:
@@ -216,7 +225,9 @@ class AgentsController:
     ) -> Agent | ExternalAgent | AssistantAgent:
         match kind:
             case AgentKind.NATIVE:
-                agent_details = parse_create_native_args(name, kind=kind, description=description, **kwargs)
+                agent_details = parse_create_native_args(
+                    name, kind=kind, description=description, **kwargs
+                )
                 agent = Agent.model_validate(agent_details)
                 AgentsController().persist_record(agent=agent, **kwargs)
             case AgentKind.EXTERNAL:
@@ -296,12 +307,17 @@ class AgentsController:
         ref_agent.collaborators = ref_collaborators
 
         return ref_agent
-    
+
     def dereference_tools(self, agent: Agent) -> Agent:
         tool_client = self.get_tool_client()
 
         deref_agent = deepcopy(agent)
-        matching_tools = tool_client.get_drafts_by_names(deref_agent.tools)
+
+        # If agent has style set to "planner" and have join_tool defined, then we need to include that tool as well
+        if agent.style == AgentStyle.PLANNER and agent.custom_join_tool:
+            matching_tools = tool_client.get_drafts_by_names(deref_agent.tools + [deref_agent.custom_join_tool])
+        else:
+            matching_tools = tool_client.get_drafts_by_names(deref_agent.tools)
 
         name_id_lookup = {}
         for tool in matching_tools:
@@ -318,6 +334,13 @@ class AgentsController:
                 sys.exit(1)
             deref_tools.append(id)
         deref_agent.tools = deref_tools
+        
+        if agent.style == AgentStyle.PLANNER and agent.custom_join_tool:
+            join_tool_id = name_id_lookup.get(agent.custom_join_tool)
+            if not join_tool_id:
+                logger.error(f"Failed to find custom join tool. No tools found with the name '{agent.custom_join_tool}'")
+                sys.exit(1)
+            deref_agent.custom_join_tool = join_tool_id
 
         return deref_agent
     
@@ -325,7 +348,12 @@ class AgentsController:
         tool_client = self.get_tool_client()
 
         ref_agent = deepcopy(agent)
-        matching_tools = tool_client.get_drafts_by_ids(ref_agent.tools)
+        
+        # If agent has style set to "planner" and have join_tool defined, then we need to include that tool as well
+        if agent.style == AgentStyle.PLANNER and agent.custom_join_tool:
+            matching_tools = tool_client.get_drafts_by_ids(ref_agent.tools + [ref_agent.custom_join_tool])
+        else:
+            matching_tools = tool_client.get_drafts_by_ids(ref_agent.tools)
 
         id_name_lookup = {}
         for tool in matching_tools:
@@ -342,6 +370,13 @@ class AgentsController:
                 sys.exit(1)
             ref_tools.append(name)
         ref_agent.tools = ref_tools
+        
+        if agent.style == AgentStyle.PLANNER and agent.custom_join_tool:
+            join_tool_name = id_name_lookup.get(agent.custom_join_tool)
+            if not join_tool_name:
+                logger.error(f"Failed to find custom join tool. No tools found with the id '{agent.custom_join_tool}'")
+                sys.exit(1)
+            ref_agent.custom_join_tool = join_tool_name
 
         return ref_agent
     
@@ -409,7 +444,7 @@ class AgentsController:
     def dereference_native_agent_dependencies(self, agent: Agent) -> Agent:
         if agent.collaborators and len(agent.collaborators):
             agent = self.dereference_collaborators(agent)
-        if agent.tools and len(agent.tools):
+        if (agent.tools and len(agent.tools)) or (agent.style == AgentStyle.PLANNER and agent.custom_join_tool):
             agent = self.dereference_tools(agent)
         if agent.knowledge_base and len(agent.knowledge_base):
             agent = self.dereference_knowledge_bases(agent)
@@ -419,7 +454,7 @@ class AgentsController:
     def reference_native_agent_dependencies(self, agent: Agent) -> Agent:
         if agent.collaborators and len(agent.collaborators):
             agent = self.reference_collaborators(agent)
-        if agent.tools and len(agent.tools):
+        if (agent.tools and len(agent.tools)) or (agent.style == AgentStyle.PLANNER and agent.custom_join_tool):
             agent = self.reference_tools(agent)
         if agent.knowledge_base and len(agent.knowledge_base):
             agent = self.reference_knowledge_bases(agent)
@@ -443,21 +478,21 @@ class AgentsController:
         return agent
     
     # Convert all names used in an agent to the corresponding ids
-    def dereference_agent_dependencies(self, agent: Agent | ExternalAgent | AssistantAgent ) -> Agent | ExternalAgent | AssistantAgent:
+    def dereference_agent_dependencies(self, agent: AnyAgentT) -> AnyAgentT:
         if isinstance(agent, Agent):
             return self.dereference_native_agent_dependencies(agent)
         if isinstance(agent, ExternalAgent) or isinstance(agent, AssistantAgent):
             return self.dereference_external_or_assistant_agent_dependencies(agent)
-    
+
     # Convert all ids used in an agent to the corresponding names
-    def reference_agent_dependencies(self, agent: Agent | ExternalAgent | AssistantAgent ) -> Agent | ExternalAgent | AssistantAgent:
+    def reference_agent_dependencies(self, agent: AnyAgentT) -> AnyAgentT:
         if isinstance(agent, Agent):
             return self.reference_native_agent_dependencies(agent)
         if isinstance(agent, ExternalAgent) or isinstance(agent, AssistantAgent):
             return self.reference_external_or_assistant_agent_dependencies(agent)
 
     def publish_or_update_agents(
-        self, agents: Iterable[Agent]
+        self, agents: Iterable[Agent | ExternalAgent | AssistantAgent]
     ):
         for agent in agents:
             agent_name = agent.name
@@ -476,6 +511,18 @@ class AgentsController:
             all_existing_agents = existing_external_clients + existing_native_agents + existing_assistant_clients
             agent = self.dereference_agent_dependencies(agent)
 
+            if isinstance(agent, Agent) and agent.style == AgentStyle.PLANNER and isinstance(agent.custom_join_tool, str):
+                tool_client = self.get_tool_client()
+
+                join_tool_spec = ToolSpec.model_validate(
+                    tool_client.get_draft_by_id(agent.custom_join_tool)
+                )
+                if not join_tool_spec.is_custom_join_tool():
+                    logger.error(
+                        f"Tool '{join_tool_spec.name}' configured as the custom join tool is not a valid join tool. A custom join tool must be a Python tool with specific input and output schema."
+                    )
+                    sys.exit(1)
+
             agent_kind = agent.kind
 
             if len(all_existing_agents) > 1:
@@ -762,32 +809,32 @@ class AgentsController:
                 rich.print(assistants_table)
 
     def remove_agent(self, name: str, kind: AgentKind):
-            try:
-                if kind == AgentKind.NATIVE:
-                    client = self.get_native_client()
-                elif kind == AgentKind.EXTERNAL:
-                    client = self.get_external_client()
-                elif kind == AgentKind.ASSISTANT:
-                    client = self.get_assistant_client()
-                else:
-                    raise ValueError("'kind' must be 'native'")
-
-                draft_agents = client.get_draft_by_name(name)
-                if len(draft_agents) > 1:
-                    logger.error(f"Multiple '{kind}' agents found with name '{name}'. Failed to delete agent")
-                    sys.exit(1)
-                if len(draft_agents) > 0:
-                    draft_agent = draft_agents[0]
-                    agent_id = draft_agent.get("id")
-                    client.delete(agent_id=agent_id)
-                    
-                    logger.info(f"Successfully removed agent {name}")
-                else:
-                    logger.warning(f"No agent named '{name}' found")
-            except requests.HTTPError as e:
-                logger.error(e.response.text)
-                exit(1)
-    
+        try:
+            if kind == AgentKind.NATIVE:
+                client = self.get_native_client()
+            elif kind == AgentKind.EXTERNAL:
+                client = self.get_external_client()
+            elif kind == AgentKind.ASSISTANT:
+                client = self.get_assistant_client()
+            else:
+                raise ValueError("'kind' must be 'native'")
+
+            draft_agents = client.get_draft_by_name(name)
+            if len(draft_agents) > 1:
+                logger.error(f"Multiple '{kind}' agents found with name '{name}'. Failed to delete agent")
+                sys.exit(1)
+            if len(draft_agents) > 0:
+                draft_agent = draft_agents[0]
+                agent_id = draft_agent.get("id")
+                client.delete(agent_id=agent_id)
+
+                logger.info(f"Successfully removed agent {name}")
+            else:
+                logger.warning(f"No agent named '{name}' found")
+        except requests.HTTPError as e:
+            logger.error(e.response.text)
+            exit(1)
+
     def get_spec_file_content(self, agent: Agent | ExternalAgent | AssistantAgent):
         ref_agent = self.reference_agent_dependencies(agent)
         agent_spec = ref_agent.model_dump(mode='json', exclude_none=True)
@@ -810,7 +857,7 @@ class AgentsController:
         
         return agent
     
-    def get_agent_by_id(self, id: str) -> Agent | ExternalAgent | AssistantAgent:
+    def get_agent_by_id(self, id: str) -> Agent | ExternalAgent | AssistantAgent | None:
         native_client = self.get_native_client()
         external_client = self.get_external_client()
         assistant_client = self.get_assistant_client()