ibm-watsonx-orchestrate 1.8.0b1__py3-none-any.whl → 1.9.0__py3-none-any.whl

ibm_watsonx_orchestrate/__init__.py

@@ -5,7 +5,7 @@
 
 pkg_name = "ibm-watsonx-orchestrate"
 
-__version__ = "1.8.0b1"
+__version__ = "1.9.0"
 
 
 

ibm_watsonx_orchestrate/agent_builder/knowledge_bases/types.py

@@ -86,11 +86,10 @@ class GenerationConfiguration(BaseModel):
     {
         "model_id": "meta-llama/llama-3-1-70b-instruct",
         "prompt_instruction": "When the documents are in different languages, you should respond in english.",
-        "retrieval_confidence_threshold": "Lowest",
         "generated_response_length": "Moderate",
-        "response_confidence_threshold": "Low",
         "display_text_no_results_found": "no docs found",
         "display_text_connectivity_issue": "conn failed",
+        "idk_message": "I dont know",
     }
     """
 
@@ -99,6 +98,7 @@ class GenerationConfiguration(BaseModel):
     generated_response_length: Optional[GeneratedResponseLength] = None
     display_text_no_results_found: Optional[str] = None
     display_text_connectivity_issue: Optional[str] = None
+    idk_message: Optional[str] = None
 
 class FieldMapping(BaseModel):
     """

ibm_watsonx_orchestrate/agent_builder/models/types.py

@@ -87,6 +87,11 @@ class ProviderConfig(BaseModel):
     azure_entra_tenant_id: Optional[str] = Field(None, alias="azureEntraTenantId")
     azure_ad_token: Optional[str] = Field(None, alias="azureAdToken")
     azure_model_name: Optional[str] = Field(None, alias="azureModelName")
+    azure_inference_deployment_name: Optional[str] = Field(None, alias="azureDeploymentName")
+    azure_inference_api_version: Optional[str] = Field(None, alias="azureApiVersion")
+    azure_inference_extra_params: Optional[str] = Field(None, alias="azureExtraParams")
+    azure_inference_foundry_url: Optional[str] = Field(None, alias="azureFoundryUrl")
+
 
     # Workers AI specific
     workers_ai_account_id: Optional[str] = Field(None, alias="workersAiAccountId")

ibm_watsonx_orchestrate/agent_builder/tools/openapi_tool.py

@@ -14,7 +14,7 @@ from ibm_watsonx_orchestrate.utils.utils import yaml_safe_load
 from .types import ToolSpec
 from .base_tool import BaseTool
 from .types import HTTP_METHOD, ToolPermission, ToolRequestBody, ToolResponseBody, \
-    OpenApiToolBinding, \
+    OpenApiToolBinding, AcknowledgementBinding, \
     JsonSchemaObject, ToolBinding, OpenApiSecurityScheme, CallbackBinding
 
 import json
@@ -207,35 +207,82 @@ def create_openapi_json_tool(
 
     # If it's an async tool, add callback binding
     if spec.is_async:
-
-
         callbacks = route_spec.get('callbacks', {})
         callback_name = next(iter(callbacks.keys()))
         callback_spec = callbacks[callback_name]
         callback_path = next(iter(callback_spec.keys()))
         callback_method = next(iter(callback_spec[callback_path].keys()))
+        callback_operation = callback_spec[callback_path][callback_method]
         
-        # Phase 1: Create a separate input schema for callback that excludes callbackUrl
-        # Note: Currently assuming the callback URL parameter will be named 'callbackUrl' in the OpenAPI spec
-        # Future phases will handle other naming conventions
+        # Extract callback input schema from the callback requestBody
         callback_input_schema = ToolRequestBody(
+            type='object',
+            properties={},
+            required=[]
+        )
+        
+        # Handle callback parameters (query, path, header params)
+        callback_parameters = callback_operation.get('parameters') or []
+        for parameter in callback_parameters:
+            name = f"{parameter['in']}_{parameter['name']}"
+            if parameter.get('required'):
+                callback_input_schema.required.append(name)
+            parameter['schema']['title'] = parameter['name']
+            parameter['schema']['description'] = parameter.get('description', None)
+            callback_input_schema.properties[name] = JsonSchemaObject.model_validate(parameter['schema'])
+            callback_input_schema.properties[name].in_field = parameter['in']
+            callback_input_schema.properties[name].aliasName = parameter['name']
+        
+        # Handle callback request body
+        callback_request_body_params = callback_operation.get('requestBody', {}).get('content', {}).get(http_response_content_type, {}).get('schema', None)
+        if callback_request_body_params is not None:
+            callback_input_schema.required.append('__requestBody__')
+            callback_request_body_params = copy.deepcopy(callback_request_body_params)
+            callback_request_body_params['in'] = 'body'
+            if callback_request_body_params.get('title') is None:
+                callback_request_body_params['title'] = 'CallbackRequestBody'
+            if callback_request_body_params.get('description') is None:
+                callback_request_body_params['description'] = 'The callback request body used for this async operation.'
+            
+            callback_input_schema.properties['__requestBody__'] = JsonSchemaObject.model_validate(callback_request_body_params)
+        
+        # Extract callback output schema
+        callback_responses = callback_operation.get('responses', {})
+        callback_response = callback_responses.get(str(http_success_response_code), {})
+        callback_response_description = callback_response.get('description')
+        callback_response_schema = callback_response.get('content', {}).get(http_response_content_type, {}).get('schema', {})
+        
+        callback_response_schema['required'] = []
+        callback_output_schema = ToolResponseBody.model_validate(callback_response_schema)
+        callback_output_schema.description = callback_response_description
+
+        # Remove callbackUrl parameter from main tool input schema
+        original_input_schema = ToolRequestBody(
             type='object',
             properties={k: v for k, v in spec.input_schema.properties.items() if not k.endswith('_callbackUrl')},
             required=[r for r in spec.input_schema.required if not r.endswith('_callbackUrl')]
         )
+        spec.input_schema = original_input_schema
 
-        if callback_input_schema:
-            spec.input_schema = callback_input_schema
-
+        original_response_schema = spec.output_schema
+        
         callback_binding = CallbackBinding(
             callback_url=callback_path,
             method=callback_method.upper(),
-            input_schema=callback_input_schema,
-            output_schema=spec.output_schema
+            output_schema=callback_output_schema
         )
 
+        # Create acknowledgement binding with the original response schema
+        acknowledgement_binding = AcknowledgementBinding(
+            output_schema=original_response_schema
+        )
+        
+        # For async tools, set the main tool's output_schema to the callback's input_schema
+        spec.output_schema = callback_input_schema
+
     else:
         callback_binding = None
+        acknowledgement_binding = None
 
     openapi_binding = OpenApiToolBinding(
         http_path=http_path,
@@ -248,6 +295,9 @@ def create_openapi_json_tool(
     if callback_binding is not None:
         openapi_binding.callback = callback_binding
 
+    if acknowledgement_binding is not None:
+        openapi_binding.acknowledgement = acknowledgement_binding
+
     spec.binding = ToolBinding(openapi=openapi_binding)
 
     return OpenAPITool(spec=spec)

ibm_watsonx_orchestrate/agent_builder/tools/python_tool.py

@@ -187,15 +187,27 @@ def _fix_optional(schema):
     replacements = {}
     if schema.required is None:
         schema.required = []
-
     for k, v in schema.properties.items():
+        # Simple null type & required -> not required
         if v.type == 'null' and k in schema.required:
             not_required.append(k)
-        if v.anyOf is not None and next(filter(lambda x: x.type == 'null', v.anyOf)) and k in schema.required:
-            v.anyOf = list(filter(lambda x: x.type != 'null', v.anyOf))
-            if len(v.anyOf) == 1:
-                replacements[k] = v.anyOf[0]
-            not_required.append(k)
+        # Optional with null & required
+        if v.anyOf is not None and [x for x in v.anyOf if x.type == 'null']:
+            if k in schema.required:
+            # required with default -> not required 
+            # required without default -> required & remove null from union
+                if v.default:
+                    not_required.append(k)
+                else:
+                    v.anyOf = list(filter(lambda x: x.type != 'null', v.anyOf))
+                if len(v.anyOf) == 1:
+                    replacements[k] = v.anyOf[0]
+            else:
+            # not required with default -> no change
+            # not required without default -> means default input is 'None'
+                v.default = v.default if v.default else 'null'
+
+
     schema.required = list(filter(lambda x: x not in not_required, schema.required if schema.required is not None else []))
     for k, v in replacements.items():
         combined = {

ibm_watsonx_orchestrate/agent_builder/tools/types.py

@@ -36,6 +36,7 @@ class JsonSchemaObject(BaseModel):
     anyOf: Optional[List['JsonSchemaObject']] = None
     in_field: Optional[Literal['query', 'header', 'path', 'body']] = Field(None, alias='in')
     aliasName: str | None = None
+    wrap_data: Optional[bool] = True
     "Runtime feature where the sdk can provide the original name of a field before prefixing"
 
     @model_validator(mode='after')
@@ -48,9 +49,9 @@ class JsonSchemaObject(BaseModel):
 class ToolRequestBody(BaseModel):
     model_config = ConfigDict(extra='allow')
 
-    type: Literal['object']
-    properties: Dict[str, JsonSchemaObject]
-    required: Optional[List[str]] = None
+    type: Literal['object', 'string']
+    properties: Optional[Dict[str, JsonSchemaObject]] = {}
+    required: Optional[List[str]] = []
 
 
 class ToolResponseBody(BaseModel):
@@ -97,9 +98,13 @@ HTTP_METHOD = Literal['GET', 'POST', 'PUT', 'PATCH', 'DELETE']
 class CallbackBinding(BaseModel):
     callback_url: str
     method: HTTP_METHOD
-    input_schema: ToolRequestBody
+    input_schema: Optional[ToolRequestBody] = None
+    output_schema: ToolResponseBody
+
+class AcknowledgementBinding(BaseModel):
     output_schema: ToolResponseBody
 
+
 class OpenApiToolBinding(BaseModel):
     http_method: HTTP_METHOD
     http_path: str
@@ -107,7 +112,8 @@ class OpenApiToolBinding(BaseModel):
     security: Optional[List[OpenApiSecurityScheme]] = None
     servers: Optional[List[str]] = None
     connection_id: str | None = None
-    callback: CallbackBinding = None
+    callback: Optional[CallbackBinding] = None
+    acknowledgement: Optional[AcknowledgementBinding] = None
 
     @model_validator(mode='after')
     def validate_openapi_tool_binding(self):
@@ -181,6 +187,7 @@ class ToolBinding(BaseModel):
 
 class ToolSpec(BaseModel):
     name: str
+    id: str | None = None
     display_name: str | None = None
     description: str
     permission: ToolPermission

ibm_watsonx_orchestrate/cli/commands/agents/agents_controller.py

@@ -750,6 +750,13 @@ class AgentsController:
     def list_agents(self, kind: AgentKind=None, verbose: bool=False):
         parse_errors = []
 
+        if verbose:
+            verbose_output_dictionary = {
+                "native":[],
+                "assistant":[],
+                "external":[]  
+            }
+
         if kind == AgentKind.NATIVE or kind is None:
             response = self.get_native_client().get()
             native_agents = []
@@ -769,7 +776,7 @@ class AgentsController:
                 for agent in native_agents:
                     agents_list.append(json.loads(agent.dumps_spec()))
 
-                rich.print(rich.json.JSON(json.dumps(agents_list, indent=4)))
+                verbose_output_dictionary["native"] = agents_list
             else:
                 native_table = rich.table.Table(
                     show_header=True, 
@@ -832,7 +839,7 @@ class AgentsController:
             if verbose:
                 for agent in external_agents:
                     external_agents_list.append(json.loads(agent.dumps_spec()))
-                rich.print(rich.json.JSON(json.dumps(external_agents_list, indent=4)))
+                verbose_output_dictionary["external"] = external_agents_list
             else:
                 external_table = rich.table.Table(
                     show_header=True, 
@@ -899,8 +906,10 @@ class AgentsController:
                     assistant_agent.config.authorization_url = response_data.get("authorization_url", assistant_agent.config.authorization_url)
 
             if verbose:
+                assistant_agent_specs = []
                 for agent in assistant_agents:
-                    rich.print(agent.dumps_spec())
+                    assistant_agent_specs.append(json.loads(agent.dumps_spec()))
+                verbose_output_dictionary["assistant"] = assistant_agent_specs
             else:
                 assistants_table = rich.table.Table(
                     show_header=True, 
@@ -938,6 +947,9 @@ class AgentsController:
                     )
                 rich.print(assistants_table)
 
+        if verbose:
+            rich.print_json(data=verbose_output_dictionary)
+
         for error in parse_errors:
             for l in error:
                 logger.error(l)

ibm_watsonx_orchestrate/cli/commands/channels/types.py

@@ -13,8 +13,8 @@ class ChannelType(str, Enum):
 
     def __str__(self):
         return self.value
-    
-    
+
+
 class RuntimeEnvironmentType(str, Enum):
     LOCAL = 'local'
     CPD = 'cpd'

ibm_watsonx_orchestrate/cli/commands/channels/webchat/channels_webchat_controller.py

@@ -11,6 +11,7 @@ from ibm_watsonx_orchestrate.client.agents.agent_client import AgentClient
 
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
 
+
 logger = logging.getLogger(__name__)
 
 class ChannelsWebchatController:
@@ -96,8 +97,6 @@ class ChannelsWebchatController:
         if target_env == 'draft' and is_saas == True:
             logger.error(f'For SAAS, please ensure this agent exists in a Live Environment')
             exit(1)
-             
-
 
         return filtered_environments[0].get("id")
 
@@ -182,7 +181,7 @@ class ChannelsWebchatController:
             case _:
                 logger.error("Environment not recognized")
                 sys.exit(1)
-            
+
         host_url = self.get_host_url()
         agent_id = self.get_agent_id(self.agent_name)
         agent_env_id = self.get_environment_id(self.agent_name, self.env)

ibm_watsonx_orchestrate/cli/commands/connections/connections_controller.py

@@ -169,7 +169,7 @@ def _validate_connection_params(type: ConnectionType, **args) -> None:
     
 
 def _parse_entry(entry: str) -> dict[str,str]:
-    split_entry = entry.split('=')
+    split_entry = entry.split('=', 1)
     if len(split_entry) != 2:
         message = f"The entry '{entry}' is not in the expected form '<key>=<value>'"
         logger.error(message)
@@ -404,8 +404,11 @@ def list_connections(environment: ConnectionEnvironment | None, verbose: bool =
                     "❌"
                 )
                 continue
-
-            connection_type = get_connection_type(security_scheme=conn.security_scheme, auth_type=conn.auth_type)
+            
+            try:
+                connection_type = get_connection_type(security_scheme=conn.security_scheme, auth_type=conn.auth_type)
+            except:
+                connection_type = conn.auth_type
 
             if conn.environment == ConnectionEnvironment.DRAFT:
                 draft_table.add_row(

ibm_watsonx_orchestrate/cli/commands/copilot/copilot_controller.py

@@ -10,10 +10,12 @@ from rich.progress import Progress, SpinnerColumn, TextColumn
 from requests import ConnectionError
 from typing import List
 from ibm_watsonx_orchestrate.client.base_api_client import ClientAPIException
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import KnowledgeBaseSpec
 from ibm_watsonx_orchestrate.agent_builder.tools import ToolSpec, ToolPermission, ToolRequestBody, ToolResponseBody
 from ibm_watsonx_orchestrate.cli.commands.agents.agents_controller import AgentsController, AgentKind, SpecVersion
 from ibm_watsonx_orchestrate.agent_builder.agents.types import DEFAULT_LLM, BaseAgentSpec
 from ibm_watsonx_orchestrate.client.agents.agent_client import AgentClient
+from ibm_watsonx_orchestrate.client.knowledge_bases.knowledge_base_client import KnowledgeBaseClient
 from ibm_watsonx_orchestrate.client.tools.tool_client import ToolClient
 from ibm_watsonx_orchestrate.client.copilot.cpe.copilot_cpe_client import CPEClient
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
@@ -56,10 +58,16 @@ def _get_incomplete_tool_from_name(tool_name: str) -> dict:
                        "input_schema": input_schema, "output_schema": output_schema})
     return spec.model_dump()
 
+
 def _get_incomplete_agent_from_name(agent_name: str) -> dict:
     spec = BaseAgentSpec(**{"name": agent_name, "description": agent_name, "kind": AgentKind.NATIVE})
     return spec.model_dump()
 
+def _get_incomplete_knowledge_base_from_name(kb_name: str) -> dict:
+    spec = KnowledgeBaseSpec(**{"name": kb_name, "description": kb_name})
+    return spec.model_dump()
+
+
 def _get_tools_from_names(tool_names: List[str]) -> List[dict]:
     if not len(tool_names):
         return []
@@ -115,6 +123,34 @@ def _get_agents_from_names(collaborators_names: List[str]) -> List[dict]:
 
     return agents
 
+def _get_knowledge_bases_from_names(kb_names: List[str]) -> List[dict]:
+    if not len(kb_names):
+        return []
+
+    kb_client = get_knowledge_bases_client()
+
+    try:
+        with _get_progress_spinner() as progress:
+            task = progress.add_task(description="Fetching Knowledge Bases", total=None)
+            knowledge_bases = kb_client.get_by_names(kb_names)
+            found_kbs = {kb.get("name") for kb in knowledge_bases}
+            progress.remove_task(task)
+            progress.refresh()
+            for kb_name in kb_names:
+                if kb_name not in found_kbs:
+                    logger.warning(
+                        f"Failed to find knowledge base named '{kb_name}'. Falling back to incomplete knowledge base definition. Copilot performance maybe effected.")
+                    knowledge_bases.append(_get_incomplete_knowledge_base_from_name(kb_name))
+    except ConnectionError:
+        logger.warning(
+            f"Failed to fetch knowledge bases from server. For optimal results please start the server and import the relevant knowledge bases {', '.join(kb_names)}.")
+        knowledge_bases = []
+        for kb_name in kb_names:
+            knowledge_bases.append(_get_incomplete_knowledge_base_from_name(kb_name))
+
+    return knowledge_bases
+
+
 def get_cpe_client() -> CPEClient:
     url = os.getenv('CPE_URL', "http://localhost:8081")
     return instantiate_client(client=CPEClient, url=url)
@@ -124,6 +160,10 @@ def get_tool_client(*args, **kwargs):
     return instantiate_client(ToolClient)
 
 
+def get_knowledge_bases_client(*args, **kwargs):
+    return instantiate_client(KnowledgeBaseClient)
+
+
 def get_native_client(*args, **kwargs):
     return instantiate_client(AgentClient)
 
@@ -137,9 +177,6 @@ def gather_utterances(max: int) -> list[str]:
     while count < max:
         utterance = Prompt.ask("  [green]>[/green]").strip()
 
-        if utterance.lower() == 'q':
-            break
-
         if utterance:
             utterances.append(utterance)
             count += 1
@@ -147,18 +184,34 @@ def gather_utterances(max: int) -> list[str]:
     return utterances
 
 
-def get_deployed_tools_agents():
+def get_knowledge_bases(client):
+    with _get_progress_spinner() as progress:
+        task = progress.add_task(description="Fetching Knowledge Bases", total=None)
+        try:
+            knowledge_bases = client.get()
+            progress.remove_task(task)
+        except ConnectionError:
+            knowledge_bases = []
+            progress.remove_task(task)
+            progress.refresh()
+            logger.warning("Failed to contact wxo server to fetch knowledge_bases. Proceeding with empty agent list")
+    return knowledge_bases
+
+
+def get_deployed_tools_agents_and_knowledge_bases():
     all_tools = find_tools_by_description(tool_client=get_tool_client(), description=None)
     # TODO: this brings only the "native" agents. Can external and assistant agents also be collaborators?
     all_agents = find_agents(agent_client=get_native_client())
-    return {"tools": all_tools, "agents": all_agents}
+    all_knowledge_bases = get_knowledge_bases(get_knowledge_bases_client())
+
+    return {"tools": all_tools, "collaborators": all_agents, "knowledge_bases": all_knowledge_bases}
 
 
 def pre_cpe_step(cpe_client):
-    tools_agents = get_deployed_tools_agents()
+    tools_agents_and_knowledge_bases = get_deployed_tools_agents_and_knowledge_bases()
     user_message = ""
     with _get_progress_spinner() as progress:
-        task = progress.add_task(description="Initilizing Prompt Engine", total=None)
+        task = progress.add_task(description="Initializing Prompt Engine", total=None)
         response = cpe_client.submit_pre_cpe_chat(user_message=user_message)
         progress.remove_task(task)
 
@@ -168,15 +221,26 @@ def pre_cpe_step(cpe_client):
             rich.print('\n🤖 Copilot: ' + response["message"])
             user_message = Prompt.ask("\n👤 You").strip()
             message_content = {"user_message": user_message}
-        elif "description" in response and response["description"]:
+        elif "description" in response and response["description"]:  # after we have a description, we pass the all tools
             res["description"] = response["description"]
-            message_content = tools_agents
-        elif "metadata" in response:
-            res["agent_name"] = response["metadata"]["agent_name"]
-            res["agent_style"] = response["metadata"]["style"]
-            res["tools"] = [t for t in tools_agents["tools"] if t["name"] in response["metadata"]["tools"]]
-            res["collaborators"] = [a for a in tools_agents["agents"] if
-                                    a["name"] in response["metadata"]["collaborators"]]
+            message_content = {"tools": tools_agents_and_knowledge_bases['tools']}
+        elif "tools" in response and response[
+            'tools'] is not None:  # after tools were selected, we pass all collaborators
+            res["tools"] = [t for t in tools_agents_and_knowledge_bases["tools"] if
+                            t["name"] in response["tools"]]
+            message_content = {"collaborators": tools_agents_and_knowledge_bases['collaborators']}
+        elif "collaborators" in response and response[
+            'collaborators'] is not None:  # after we have collaborators, we pass all knowledge bases
+            res["collaborators"] = [a for a in tools_agents_and_knowledge_bases["collaborators"] if
+                                    a["name"] in response["collaborators"]]
+            message_content = {"knowledge_bases": tools_agents_and_knowledge_bases['knowledge_bases']}
+        elif "knowledge_bases" in response and response['knowledge_bases'] is not None:  # after we have knowledge bases, we pass selected=True to mark that all selection were done
+            res["knowledge_bases"] = [a for a in tools_agents_and_knowledge_bases["knowledge_bases"] if
+                                      a["name"] in response["knowledge_bases"]]
+            message_content = {"selected": True}
+        elif "agent_name" in response and response['agent_name'] is not None:  # once we have a name and style, this phase has ended
+            res["agent_name"] = response["agent_name"]
+            res["agent_style"] = response["agent_style"]
             return res
         with _get_progress_spinner() as progress:
             task = progress.add_task(description="Thinking...", total=None)
@@ -197,6 +261,7 @@ def find_tools_by_description(description, tool_client):
             logger.warning("Failed to contact wxo server to fetch tools. Proceeding with empty tool list")
     return tools
 
+
 def find_agents(agent_client):
     with _get_progress_spinner() as progress:
         task = progress.add_task(description="Fetching Agents", total=None)
@@ -282,16 +347,25 @@ def prompt_tune(agent_spec: str, output_file: str | None, samples_file: str | No
     tools = _get_tools_from_names(agent.tools)
 
     collaborators = _get_agents_from_names(agent.collaborators)
+
+    knowledge_bases = _get_knowledge_bases_from_names(agent.knowledge_base)
     try:
-        new_prompt = talk_to_cpe(cpe_client=client, samples_file=samples_file,
-                                 context_data={"initial_instruction": instr, 'tools': tools, 'description': agent.description,
-                                               "collaborators": collaborators})
+        new_prompt = talk_to_cpe(cpe_client=client,
+                                samples_file=samples_file,
+                                context_data={
+                                    "initial_instruction": instr,
+                                    'tools': tools,
+                                    'description': agent.description,
+                                    "collaborators": collaborators,
+                                    "knowledge_bases": knowledge_bases
+                                })
     except ConnectionError:
         logger.error(
             "Failed to connect to Copilot server. Please ensure Copilot is running via `orchestrate copilot start`")
         sys.exit(1)
     except ClientAPIException:
-        logger.error("An unexpected server error has occur with in the Copilot server. Please check the logs via `orchestrate server logs`")
+        logger.error(
+            "An unexpected server error has occur with in the Copilot server. Please check the logs via `orchestrate server logs`")
         sys.exit(1)
 
     if new_prompt:
@@ -319,17 +393,21 @@ def create_agent(output_file: str, llm: str, samples_file: str | None, dry_run_f
             "Failed to connect to Copilot server. Please ensure Copilot is running via `orchestrate copilot start`")
         sys.exit(1)
     except ClientAPIException:
-        logger.error("An unexpected server error has occur with in the Copilot server. Please check the logs via `orchestrate server logs`")
+        logger.error(
+            "An unexpected server error has occur with in the Copilot server. Please check the logs via `orchestrate server logs`")
         sys.exit(1)
-        
+
     tools = res["tools"]
     collaborators = res["collaborators"]
+    knowledge_bases = res["knowledge_bases"]
     description = res["description"]
     agent_name = res["agent_name"]
     agent_style = res["agent_style"]
 
     # 4. discuss the instructions
-    instructions = talk_to_cpe(cpe_client, samples_file, {'description': description, 'tools': tools, 'collaborators': collaborators})
+    instructions = talk_to_cpe(cpe_client, samples_file,
+                               {'description': description, 'tools': tools, 'collaborators': collaborators,
+                                'knowledge_bases': knowledge_bases})
 
     # 6. create and save the agent
     llm = llm if llm else DEFAULT_LLM
@@ -337,7 +415,9 @@ def create_agent(output_file: str, llm: str, samples_file: str | None, dry_run_f
         'style': agent_style,
         'tools': [t['name'] for t in tools],
         'llm': llm,
-        'collaborators': [c['name'] for c in collaborators]
+        'collaborators': [c['name'] for c in collaborators],
+        'knowledge_base': [k['name'] for k in knowledge_bases]
+        # generate_agent_spec expects knowledge_base and not knowledge_bases
     }
     agent = AgentsController.generate_agent_spec(agent_name, AgentKind.NATIVE, description, **params)
     agent.instructions = instructions