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

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()

ibm_watsonx_orchestrate/cli/commands/models/model_provider_mapper.py

@@ -0,0 +1,191 @@
+# https://portkey.ai/
+from dotenv import dotenv_values
+import sys
+import logging
+from ibm_watsonx_orchestrate.agent_builder.models.types import  ProviderConfig, ModelProvider
+
+logger = logging.getLogger(__name__)
+
+_BASIC_PROVIDER_CONFIG_KEYS = {'provider', 'api_key', 'custom_host', 'url_to_fetch', 'forward_headers', 'request_timeout', 'transform_to_form_data'}
+
+PROVIDER_EXTRA_PROPERTIES_LUT = {
+    ModelProvider.ANTHROPIC: {'anthropic_beta', 'anthropic_version'},
+    # ModelProvider.AZURE_AI: {
+    #     'azure_resource_name',
+    #     'azure_deployment_id',
+    #     'azure_api_version',
+    #     'ad_auth',
+    #     'azure_auth_mode',
+    #     'azure_managed_client_id',
+    #     'azure_entra_client_id',
+    #     'azure_entra_client_secret',
+    #     'azure_entra_tenant_id',
+    #     'azure_ad_token',
+    #     'azure_model_name'
+    # },
+    # ModelProvider.BEDROCK: {
+    #     'aws_secret_access_key',
+    #     'aws_access_key_id',
+    #     'aws_session_token',
+    #     'aws_region',
+    #     'aws_auth_type',
+    #     'aws_role_arn',
+    #     'aws_external_id',
+    #     'aws_s3_bucket',
+    #     'aws_s3_object_key',
+    #     'aws_bedrock_model',
+    #     'aws_server_side_encryption',
+    #     'aws_server_side_encryption_kms_key_id'
+    # },
+    # ModelProvider.VERTEX_AI: {
+    #     'vertex_region',
+    #     'vertex_project_id',
+    #     'vertex_service_account_json',
+    #     'vertex_storage_bucket_name',
+    #     'vertex_model_name',
+    #     'filename'
+    # },
+    # ModelProvider.HUGGINGFACE: {'huggingfaceBaseUrl'},
+    ModelProvider.MISTRAL_AI: {'mistral_fim_completion'},
+    # ModelProvider.STABILITY_AI: {'stability_client_id', 'stability_client_user_id', 'stability_client_version'},
+    ModelProvider.WATSONX: {'watsonx_version', 'watsonx_space_id', 'watsonx_project_id', 'api_key', 'watsonx_deployment_id', 'watsonx_cpd_url', 'watsonx_cpd_username', 'watsonx_cpd_password'},
+
+    # 'palm': _bpp('PALM', {}),
+    # 'nomic': _bpp('NOMIC', {}),
+    # 'perplexity-ai': _bpp('PERPLEXITY_AI', {}),
+    # 'segmind': _bpp('SEGMIND', {}),
+    # 'deepinfra': _bpp('DEEPINFRA', {}),
+    # 'novita-ai': _bpp('NOVITA_AI', {}),
+    # 'fireworks-ai': _bpp('FIREWORKS',{
+    #     'FIREWORKS_ACCOUNT_ID': 'fireworks_account_id'
+    # }),
+    # 'deepseek': _bpp('DEEPSEEK', {}),
+    # 'voyage': _bpp('VOYAGE', {}),
+    # 'moonshot': _bpp('MOONSHOT', {}),
+    # 'lingyi': _bpp('LINGYI', {}),
+    # 'zhipu': _bpp('ZHIPU', {}),
+    # 'monsterapi': _bpp('MONSTERAPI', {}),
+    # 'predibase': _bpp('PREDIBASE', {}),
+
+    # 'github': _bpp('GITHUB', {}),
+    # 'deepbricks': _bpp('DEEPBRICKS', {}),
+    # 'sagemaker': _bpp('AMZN_SAGEMAKER', {
+    #     'AMZN_SAGEMAKER_CUSTOM_ATTRIBUTES': 'amzn_sagemaker_custom_attributes',
+    #     'AMZN_SAGEMAKER_TARGET_MODEL': 'amzn_sagemaker_target_model',
+    #     'AMZN_SAGEMAKER_TARGET_VARIANT': 'amzn_sagemaker_target_variant',
+    #     'AMZN_SAGEMAKER_TARGET_CONTAINER_HOSTNAME': 'amzn_sagemaker_target_container_hostname',
+    #     'AMZN_SAGEMAKER_INFERENCE_ID': 'amzn_sagemaker_inference_id',
+    #     'AMZN_SAGEMAKER_ENABLE_EXPLANATIONS': 'amzn_sagemaker_enable_explanations',
+    #     'AMZN_SAGEMAKER_INFERENCE_COMPONENT': 'amzn_sagemaker_inference_component',
+    #     'AMZN_SAGEMAKER_SESSION_ID': 'amzn_sagemaker_session_id',
+    #     'AMZN_SAGEMAKER_MODEL_NAME': 'amzn_sagemaker_model_name'
+    # }),
+    # '@cf': _bpp('WORKERS_AI', { # workers ai
+    #     'WORKERS_AI_ACCOUNT_ID': 'workers_ai_account_id'
+    # }),
+    # 'snowflake': _bpp('SNOWFLAKE', { # no provider prefix found
+    #     'SNOWFLAKE_ACCOUNT': 'snowflake_account'
+    # })
+}
+
+PROVIDER_REQUIRED_FIELDS = {k:['api_key'] for k in ModelProvider}
+# Update required fields for each provider
+# Use sets to denote when a requirement is 'or'
+PROVIDER_REQUIRED_FIELDS.update({
+    ModelProvider.WATSONX: PROVIDER_REQUIRED_FIELDS[ModelProvider.WATSONX] + [{'watsonx_space_id', 'watsonx_project_id', 'watsonx_deployment_id'}],
+    ModelProvider.OLLAMA: PROVIDER_REQUIRED_FIELDS[ModelProvider.OLLAMA] + ['custom_host']
+})
+
+# def env_file_to_model_ProviderConfig(model_name: str, env_file_path: str) -> ProviderConfig | None:
+#     provider = next(filter(lambda x: x not in ('virtual-policy', 'virtual-model'), model_name.split('/')))
+#     if provider not in ModelProvider:
+#         logger.error(f"Unsupported model provider {provider}")
+#         sys.exit(1)
+
+#     values = dotenv_values(str(env_file_path))
+
+#     if values is None:
+#         logger.error(f"No provider configuration in env file {env_file_path}")
+#         sys.exit(1)
+
+#     cfg = ProviderConfig()
+#     cfg.provider = PROVIDER_LUT[provider]
+
+#     cred_lut = PROVIDER_PROPERTIES_LUT[provider]
+
+
+#     consumed_credentials = []
+#     # Ollama requires some apikey but its content don't matter
+#     # Default it to 'ollama' to avoid users needing to specify
+#     if cfg.provider == ModelProvider.OLLAMA:
+#         consumed_credentials.append('api_key')
+#         setattr(cfg, 'api_key', ModelProvider.OLLAMA)
+    
+#     for key, value in values.items():
+#         if key in cred_lut:
+#             k = cred_lut[key]
+#             consumed_credentials.append(k)
+#             setattr(cfg, k, value)
+
+#     return cfg
+
+def _validate_provider(provider: str | ModelProvider) -> None:
+    if not ModelProvider.has_value(provider):
+        logger.error(f"Unsupported model provider {provider}")
+        sys.exit(1)
+
+def _validate_extra_fields(provider: ModelProvider, cfg: ProviderConfig) -> None:
+    accepted_fields = _BASIC_PROVIDER_CONFIG_KEYS.copy()
+    extra_accepted_fields = PROVIDER_EXTRA_PROPERTIES_LUT.get(provider)
+    if extra_accepted_fields:
+        accepted_fields = accepted_fields.union(extra_accepted_fields)
+    
+    for attr in cfg.__dict__:
+        if attr.startswith("__"):
+            continue
+
+        if  cfg.__dict__.get(attr) is not None and attr not in accepted_fields:
+            logger.warning(f"The config option '{attr}' is not used by provider '{provider}'")
+
+def _validate_requirements(provider: ModelProvider, cfg: ProviderConfig, app_id: str = None) -> None:
+    provided_credentials = set([k for k,v in dict(cfg).items() if v is not None])
+    required_creds = PROVIDER_REQUIRED_FIELDS[provider]
+    missing_credentials = []
+    for cred in required_creds:
+        if isinstance(cred, set):
+            if not any(c in provided_credentials for c in cred):
+                missing_credentials.append(cred)
+        else:
+            if cred not in provided_credentials:
+                missing_credentials.append(cred)
+
+    if len(missing_credentials) > 0:
+        if not app_id:
+            missing_credentials_string = f"Missing configuration variable(s) required for the provider {provider}:"
+        else:
+            missing_credentials_string = f"The following configuration variable(s) for the provider {provider} are not in the spec provider_config:"
+        for cred in missing_credentials:
+            if isinstance(cred, set):
+                cred_str = ' or '.join(list(cred))
+            else:
+                cred_str = cred
+            missing_credentials_string += f"\n\t  - {cred_str}"
+        
+        if not app_id:
+            logger.error(missing_credentials_string)
+            logger.error("Please provide the above values in the provider config. For secret values (e.g. 'api_key') create a key_value connection `orchestrate connections add` then bind it to the model with `--app-id`")
+            sys.exit(1)
+        else:
+            logger.info(missing_credentials_string)
+            logger.info(f"Please ensure these values are set in the connection '{app_id}'.")
+
+
+def validate_ProviderConfig(cfg: ProviderConfig, app_id: str)-> None:
+    if not cfg:
+        return
+
+    provider = cfg.provider
+
+    _validate_provider(provider)
+    _validate_extra_fields(provider=provider, cfg=cfg)
+    _validate_requirements(provider=provider, cfg=cfg, app_id=app_id)