ibm-watsonx-orchestrate 1.12.0b1__py3-none-any.whl → 1.13.0b0__py3-none-any.whl

ibm_watsonx_orchestrate/cli/commands/copilot/copilot_command.py

@@ -1,7 +1,8 @@
 import typer
 from typing_extensions import Annotated
 from pathlib import Path
-from ibm_watsonx_orchestrate.cli.commands.copilot.copilot_controller import prompt_tune, create_agent
+from ibm_watsonx_orchestrate.cli.commands.copilot.copilot_controller import prompt_tune, create_agent, \
+    refine_agent_with_trajectories
 from ibm_watsonx_orchestrate.cli.commands.copilot.copilot_server_controller import start_server, stop_server
 
 copilot_app = typer.Typer(no_args_is_help=True)
@@ -62,4 +63,26 @@ def prompt_tume_command(
             samples_file=samples,
             output_file=output_file,
             dry_run_flag=dry_run_flag,
-        )
+        )
+
+@copilot_app.command(name="autotune", help="Autotune the agent's instructions by incorporating insights from chat interactions and user feedback")
+def agent_refine(
+    agent_name: Annotated[
+        str,
+        typer.Option("--agent-name", "-n", help="The name of the agent to tune"),
+    ],
+    output_file: Annotated[
+        str,
+        typer.Option("--output-file", "-o", help="Optional output file to avoid overwriting existing agent spec"),
+    ] = None,
+    use_last_chat: Annotated[
+        bool,
+        typer.Option("--use_last_chat", "-l", help="Tuning by using the last conversation with the agent instead of prompting the user to choose chats"),
+    ] = False,
+    dry_run_flag: Annotated[
+        bool,
+        typer.Option("--dry-run",
+                     help="Dry run will prevent the tuned content being saved and output the results to console"),
+    ] = False,
+):
+    refine_agent_with_trajectories(agent_name, output_file, use_last_chat, dry_run_flag)

ibm_watsonx_orchestrate/cli/commands/copilot/copilot_controller.py

@@ -2,13 +2,17 @@ import logging
 import os
 import sys
 import csv
+import difflib
+from datetime import datetime
 
 import rich
 from rich.console import Console
 from rich.prompt import Prompt
 from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.panel import Panel
+from rich.table import Table
 from requests import ConnectionError
-from typing import List
+from typing import List, Dict
 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
@@ -16,6 +20,7 @@ from ibm_watsonx_orchestrate.cli.commands.agents.agents_controller import Agents
 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.threads.threads_client import ThreadsClient
 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
@@ -63,6 +68,7 @@ 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()
@@ -123,6 +129,7 @@ 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 []
@@ -168,6 +175,10 @@ def get_native_client(*args, **kwargs):
     return instantiate_client(AgentClient)
 
 
+def get_threads_client():
+    return instantiate_client(ThreadsClient)
+
+
 def gather_utterances(max: int) -> list[str]:
     utterances = []
     logger.info("Please provide 3 sample utterances you expect your agent to handle:")
@@ -221,7 +232,8 @@ 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"]:  # after we have a description, we pass the all tools
+        elif "description" in response and response[
+            "description"]:  # after we have a description, we pass the all tools
             res["description"] = response["description"]
             message_content = {"tools": tools_agents_and_knowledge_bases['tools']}
         elif "tools" in response and response[
@@ -234,11 +246,13 @@ def pre_cpe_step(cpe_client):
             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
+        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
+        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
@@ -351,21 +365,21 @@ def prompt_tune(agent_spec: str, output_file: str | None, samples_file: str | No
     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,
-                                    "knowledge_bases": knowledge_bases
-                                })
+                                 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`")
+            "An unexpected server error has occurred with in the Copilot server. Please check the logs via `orchestrate server logs`")
         sys.exit(1)
 
     if new_prompt:
@@ -394,7 +408,7 @@ def create_agent(output_file: str, llm: str, samples_file: str | None, dry_run_f
         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`")
+            "An unexpected server error has occurred with in the Copilot server. Please check the logs via `orchestrate server logs`")
         sys.exit(1)
 
     tools = res["tools"]
@@ -446,3 +460,224 @@ def create_agent(output_file: str, llm: str, samples_file: str | None, dry_run_f
     for line in message_lines:
         rich.print("║  " + line.ljust(max_length) + "  ║")
     rich.print("╚" + "═" * frame_width + "╝")
+
+
+def _format_thread_messages(messages:List[dict]) -> List[dict]:
+    """
+        restructure and keep only the content relevant for refining the agent before sending to the refinement process
+    :param messages: List of messages as returned from the threads endpoint
+    :param messages:
+    :return: List of dictionaries where each dictionary represents a message
+    """
+    new_messages = []
+    for m in messages:
+        m_dict = {'role': m['role'], 'content': m['content'][0]['text'], 'type': 'text'} # text message
+        if m['step_history']:
+            step_history = m['step_history']
+            for step in step_history:
+                step_details = step['step_details'][0]
+                if step_details['type'] == 'tool_calls':  # tool call
+                    for t in step_details['tool_calls']:
+                        new_messages.append(
+                            {'role': m['role'], 'type': 'tool_call', 'args': t['args'], 'name': t['name']})
+                elif step_details['type'] == 'tool_response':  # tool response
+                    new_messages.append({'role': m['role'], 'type': 'tool_response', 'content': step_details['content']})
+        new_messages.append(m_dict)
+        if m['message_state']:
+            new_messages.append({'feedback': m['message_state']['content']['1']['feedback']})
+    return new_messages
+
+
+def _suggest_sorted(user_input: str, options: List[str]) -> List[str]:
+    # Sort by similarity score
+    return sorted(options, key=lambda x: difflib.SequenceMatcher(None, user_input, x).ratio(), reverse=True)
+
+
+def refine_agent_with_trajectories(agent_name: str, output_file: str | None, use_last_chat: bool=False, dry_run_flag: bool = False) -> None:
+    """
+    Refines an existing agent's instructions using user selected chat trajectories and saves the updated agent configuration.
+
+    This function performs a multi-step process to enhance an agent's prompt instructions based on user interactions:
+
+    1. **Validation**: Ensures the output file path is valid and checks if the specified agent exists. If not found,
+       it suggests similar agent names.
+    2. **Chat Retrieval**: Fetches the 10 most recent chat threads associated with the agent. If no chats are found,
+       the user is prompted to initiate a conversation.
+    3. **User Selection**: Displays a summary of recent chats and allows the user to select which ones to use for refinement.
+    4. **Refinement**: Sends selected chat messages to the Copilot Prompt Engine (CPE) to generate refined instructions.
+    5. **Update and Save**: Updates the agent's instructions and either prints the
+       updated agent (if `dry_run_flag` is True) or saves it to the specified output file.
+
+    Parameters:
+        agent_name (str): The name of the agent to refine.
+        output_file (str): Path to the file where the refined agent configuration will be saved.
+        use_last_chat(bool): If true, optimize by using the last conversation with the agent, otherwise let the use choose
+        dry_run_flag (bool): If True, prints the refined agent configuration without saving it to disk.
+
+    Returns:
+        None
+    """
+
+    _validate_output_file(output_file, dry_run_flag)
+    agents_controller = AgentsController()
+    agents_client = get_native_client()
+    threads_client = get_threads_client()
+    all_agents = agents_controller.get_all_agents(client=agents_client)
+
+    # Step 1 - validate agent exist. If not - list the agents sorted by their distance from the user input name
+    agent_id = all_agents.get(agent_name)
+    if agent_id is None:
+        if len(all_agents) == 0:
+            raise BadRequest("No agents in workspace\nCreate your first agent using `orchestrate copilot prompt-tune`")
+        else:
+            available_sorted_str = "\n".join(_suggest_sorted(agent_name, all_agents.keys()))
+            raise BadRequest(f'Agent "{agent_name}" does not exist.\n\n'
+                             f'Available agents:\n'
+                             f'{available_sorted_str}')
+
+        rich.print(Panel(message, title="Agent Lookup", border_style="blue"))
+        return
+
+    cpe_client = get_cpe_client()
+    # Step 2 - retrieve chats (threads)
+    try:
+        with _get_progress_spinner() as progress:
+            task = progress.add_task(description="Retrieve chats", total=None)
+            all_threads = threads_client.get_all_threads(agent_id)
+            if len(all_threads) == 0:
+                progress.remove_task(task)
+                progress.refresh()
+                raise BadRequest(
+                    f"No chats found for agent '{agent_name}'. To use autotune, please initiate at least one conversation with the agent. You can start a chat using `orchestrate chat start`.",
+                   )
+                return
+            last_10_threads = all_threads[:10] #TODO use batching when server allows
+            last_10_chats = [_format_thread_messages(chat) for chat in
+                             threads_client.get_threads_messages([thread['id'] for thread in last_10_threads])]
+
+            progress.remove_task(task)
+            progress.refresh()
+    except ConnectionError:
+        logger.error(
+            f"Failed to retrieve threads (chats) for agent {agent_name}")
+        sys.exit(1)
+    except ClientAPIException:
+        logger.error(
+            f"An unexpected server error has occurred while retrieving threads for agent {agent_name}. Please check the logs via `orchestrate server logs`")
+        sys.exit(1)
+
+    # Step 3 - show chats and let the user choose
+    if use_last_chat:
+        title = "Selected chat"
+    else:
+        title = "10 Most Recent Chats"
+    table = Table(title=title)
+    table.add_column("Number", justify="right")
+    table.add_column("Chat Date", justify="left")
+    table.add_column("Title", justify="left")
+    table.add_column("Last User Message", justify="left")
+    table.add_column("Last User Feedback", justify="left")
+
+    for i, (thread, chat) in enumerate(zip(last_10_threads, last_10_chats), start=1):
+        all_user_messages = [msg for msg in chat if 'role' in msg and msg['role'] == 'user']
+
+        if len(all_user_messages) == 0:
+            last_user_message = ""
+        else:
+            last_user_message = all_user_messages[-1]['content']
+        all_feedbacks = [msg for msg in chat if 'feedback' in msg and 'text' in msg['feedback']]
+        if len(all_feedbacks) == 0:
+            last_feedback = ""
+        else:
+            last_feedback = f"{'👍' if all_feedbacks[-1]['feedback']['is_positive'] else '👎'} {all_feedbacks[-1]['feedback']['text']}"
+
+        table.add_row(str(i), datetime.strptime(thread['created_on'], '%Y-%m-%dT%H:%M:%S.%fZ').strftime(
+            '%B %d, %Y at %I:%M %p'), thread['title'], last_user_message, last_feedback)
+        table.add_row("", "", "")
+        if  use_last_chat:
+            break
+
+    rich.print(table)
+
+    if use_last_chat:
+        rich.print("Tuning using the last conversation with the agent")
+        threads_messages = [last_10_chats[0]]
+    else:
+        threads_messages = get_user_selection(last_10_chats)
+
+    # Step 4 - run the refiner
+    try:
+        with _get_progress_spinner() as progress:
+            agent = agents_controller.get_agent_by_id(id=agent_id)
+            task = progress.add_task(description="Running Prompt Refiner", total=None)
+            tools_client = get_tool_client()
+            knowledge_base_client = get_knowledge_bases_client()
+            # loaded agent contains the ids of the tools/collabs/knowledge bases, convert them back to names.
+            agent.tools = [tools_client.get_draft_by_id(id)['name'] for id in agent.tools]
+            agent.knowledge_base = [knowledge_base_client.get_by_id(id)['name'] for id in agent.knowledge_base]
+            agent.collaborators = [agents_client.get_draft_by_id(id)['name'] for id in agent.collaborators]
+            tools = _get_tools_from_names(agent.tools)
+            collaborators = _get_agents_from_names(agent.collaborators)
+            knowledge_bases = _get_knowledge_bases_from_names(agent.knowledge_base)
+            if agent.instructions is None:
+                raise BadRequest("Agent must have instructions in order to use the autotune command. To build an instruction use `orchestrate copilot prompt-tune -f <path_to_agent_yaml> -o <path_to_new_agent_yaml>`")
+            response = cpe_client.refine_agent_with_chats(agent.instructions, tools=tools, collaborators=collaborators,
+                                                          knowledge_bases=knowledge_bases, trajectories_with_feedback=threads_messages)
+            progress.remove_task(task)
+            progress.refresh()
+    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 occurred with in the Copilot server. Please check the logs via `orchestrate server logs`")
+        sys.exit(1)
+
+    # Step 5 - update the agent and print/save the results
+    agent.instructions = response['instruction']
+
+    if dry_run_flag:
+        rich.print(agent.model_dump(exclude_none=True))
+        return
+
+    if os.path.dirname(output_file):
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
+    agent.id = None # remove existing agent id before saving
+    AgentsController.persist_record(agent, output_file=output_file)
+
+    logger.info(f"Your agent refinement session finished successfully!")
+    logger.info(f"Agent YAML with the updated instruction saved in file: {os.path.abspath(output_file)}")
+
+
+
+def get_user_selection(chats: List[List[Dict]]) -> List[List[Dict]]:
+    """
+    Prompts the user to select up to 5 chat threads by entering their indices.
+
+    Parameters:
+        chats (List[List[Dict]]): A list of chat threads, where each thread is a list of message dictionaries.
+
+    Returns:
+        List[List[Dict]]: A list of selected chat threads based on user input.
+    """
+    while True:
+        try:
+            eg_str = "1" if len(chats) < 2 else "1, 2"
+            input_str = input(
+                f"Please enter up to 5 indices of chats you'd like to select, separated by commas (e.g. {eg_str}): "
+            )
+
+            choices = [int(choice.strip()) for choice in input_str.split(',')]
+
+            if len(choices) > 5:
+                rich.print("You can select up to 5 chats only. Please try again.")
+                continue
+
+            if all(1 <= choice <= len(chats) for choice in choices):
+                selected_threads = [chats[choice - 1] for choice in choices]
+                return selected_threads
+            else:
+                rich.print(f"Please enter only numbers between 1 and {len(chats)}.")
+        except ValueError:
+            rich.print("Invalid input. Please enter valid integers separated by commas.")

ibm_watsonx_orchestrate/cli/commands/copilot/copilot_server_controller.py

@@ -1,9 +1,9 @@
 import logging
 import sys
-from pathlib import Path
 import time
+from pathlib import Path
+
 import requests
-from urllib.parse import urlparse
 
 from ibm_watsonx_orchestrate.cli.config import Config
 from ibm_watsonx_orchestrate.utils.docker_utils import DockerLoginService, DockerComposeCore, DockerUtils
@@ -45,7 +45,7 @@ def run_compose_lite_cpe(user_env_file: Path) -> bool:
 
     final_env_file = env_service.write_merged_env_file(merged_env_dict)
 
-    compose_core = DockerComposeCore(env_service)
+    compose_core = DockerComposeCore(env_service=env_service)
 
     result = compose_core.service_up(service_name="cpe", friendly_name="Copilot", final_env_file=final_env_file)
 
@@ -75,7 +75,7 @@ def run_compose_lite_cpe_down(is_reset: bool = False) -> None:
 
     cli_config = Config()
     env_service = EnvService(cli_config)
-    compose_core = DockerComposeCore(env_service)
+    compose_core = DockerComposeCore(env_service=env_service)
 
     result = compose_core.service_down(service_name="cpe", friendly_name="Copilot", final_env_file=final_env_file, is_reset=is_reset)
 

ibm_watsonx_orchestrate/cli/commands/environment/environment_command.py

@@ -43,9 +43,13 @@ def activate_env(
         test_package_version_override: Annotated[
             str,
             typer.Option("--test-package-version-override", help="Which prereleased package version to reference when using --registry testpypi", hidden=True),
+        ] = None,
+        skip_version_check: Annotated[
+            bool,
+            typer.Option('--skip-version-check/--enable-version-check', help='Use this flag to skip validating that adk version in use exists in pypi (for clients who mirror the ADK to a local registry and do not have local access to pypi).')
         ] = None
 ):
-    environment_controller.activate(name=name, apikey=apikey,username=username, password=password, registry=registry, test_package_version_override=test_package_version_override)
+    environment_controller.activate(name=name, apikey=apikey, username=username, password=password, registry=registry, test_package_version_override=test_package_version_override, skip_version_check=skip_version_check)
 
 
 @environment_app.command(name="add")

ibm_watsonx_orchestrate/cli/commands/environment/environment_controller.py

@@ -18,8 +18,9 @@ from ibm_watsonx_orchestrate.cli.config import (
     ENV_IAM_URL_OPT,
     ENVIRONMENTS_SECTION_HEADER,
     PROTECTED_ENV_NAME,
-    ENV_AUTH_TYPE, PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT, BYPASS_SSL, VERIFY,
-    DEFAULT_CONFIG_FILE_CONTENT
+    ENV_AUTH_TYPE, PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT,
+    BYPASS_SSL, VERIFY,
+    DEFAULT_CONFIG_FILE_CONTENT, PYTHON_REGISTRY_SKIP_VERSION_CHECK_OPT
 )
 from ibm_watsonx_orchestrate.client.client import Client
 from ibm_watsonx_orchestrate.client.client_errors import ClientError
@@ -131,7 +132,7 @@ def _login(name: str, apikey: str = None, username: str = None, password: str =
     except ClientError as e:
         raise ClientError(e)
 
-def activate(name: str, apikey: str=None, username: str=None, password: str=None, registry: RegistryType=None, test_package_version_override=None) -> None:
+def activate(name: str, apikey: str=None, username: str=None, password: str=None, registry: RegistryType=None, test_package_version_override=None, skip_version_check=None) -> None:
     cfg = Config()
     auth_cfg = Config(AUTH_CONFIG_FILE_FOLDER, AUTH_CONFIG_FILE)
     env_cfg = cfg.read(ENVIRONMENTS_SECTION_HEADER, name)
@@ -159,6 +160,8 @@ def activate(name: str, apikey: str=None, username: str=None, password: str=None
         elif cfg.read(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT) is None:
             cfg.write(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT, DEFAULT_CONFIG_FILE_CONTENT[PYTHON_REGISTRY_HEADER][PYTHON_REGISTRY_TYPE_OPT])
             cfg.write(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT, test_package_version_override)
+        if skip_version_check is not None:
+            cfg.write(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_SKIP_VERSION_CHECK_OPT, skip_version_check)
 
     logger.info(f"Environment '{name}' is now active")
     is_cpd = is_cpd_env(url)

ibm_watsonx_orchestrate/cli/commands/evaluations/evaluations_command.py

@@ -54,7 +54,8 @@ def read_env_file(env_path: Path|str) -> dict:
 def validate_watsonx_credentials(user_env_file: str) -> bool:
     required_sets = [
         ["WATSONX_SPACE_ID", "WATSONX_APIKEY"],
-        ["WO_INSTANCE", "WO_API_KEY"]
+        ["WO_INSTANCE", "WO_API_KEY"],
+        ["WO_INSTANCE", "WO_PASSWORD", "WO_USERNAME"]
     ]
     
     def has_valid_keys(env: dict) -> bool:
@@ -75,7 +76,7 @@ def validate_watsonx_credentials(user_env_file: str) -> bool:
     user_env = read_env_file(user_env_file)
     
     if not has_valid_keys(user_env):
-        logger.error("Error: The environment file does not contain the required keys: either WATSONX_SPACE_ID and WATSONX_APIKEY or WO_INSTANCE and WO_API_KEY.")
+        logger.error("Error: The environment file does not contain the required keys: either WATSONX_SPACE_ID and WATSONX_APIKEY or WO_INSTANCE and WO_API_KEY or WO_INSTANCE and WO_USERNAME and WO_PASSWORD.")
         sys.exit(1)
 
     # Update os.environ with whichever set is present

ibm_watsonx_orchestrate/cli/commands/evaluations/evaluations_controller.py

@@ -52,7 +52,7 @@ class EvaluationsController:
 
         if "WATSONX_SPACE_ID" in os.environ and "WATSONX_APIKEY" in os.environ:
             provider = "watsonx"
-        elif "WO_INSTANCE" in os.environ and "WO_API_KEY" in os.environ:
+        elif "WO_INSTANCE" in os.environ and ("WO_API_KEY" in os.environ or "WO_PASSWORD" in os.environ):
             provider = "model_proxy"
         else:
             logger.error(

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_controller.py

@@ -6,14 +6,15 @@ import logging
 import importlib
 import inspect
 from pathlib import Path
-from typing import List
+from typing import List, Any
 
 from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.knowledge_base import KnowledgeBase
 from ibm_watsonx_orchestrate.client.knowledge_bases.knowledge_base_client import KnowledgeBaseClient
 from ibm_watsonx_orchestrate.client.base_api_client import ClientAPIException
 from ibm_watsonx_orchestrate.client.connections import get_connections_client
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
-from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import FileUpload
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import FileUpload, KnowledgeBaseListEntry
+from ibm_watsonx_orchestrate.cli.common import ListFormats, rich_table_to_markdown
 
 logger = logging.getLogger(__name__)
 
@@ -198,8 +199,7 @@ class KnowledgeBaseController:
 
         logger.info(f"Knowledge base '{kb.name}' updated successfully")
 
-
-    def knowledge_base_status( self, id: str, name: str) -> None:
+    def knowledge_base_status( self, id: str, name: str, format: ListFormats = None) ->  dict | str | None:
         knowledge_base_id = self.get_id(id, name)
         response = self.get_client().status(knowledge_base_id)
 
@@ -219,13 +219,25 @@ class KnowledgeBaseController:
 
             response["id"] = kbID
         
+        if format == ListFormats.JSON:
+            return response
+        
+
         [table.add_column(to_column_name(col), {}) for col in response.keys()]
         table.add_row(*[str(val) for val in response.values()])
         
+        if format == ListFormats.Table:
+            return rich_table_to_markdown(table)
+
         rich.print(table)
 
 
-    def list_knowledge_bases(self, verbose: bool=False):
+    def list_knowledge_bases(self, verbose: bool=False, format: ListFormats=None)-> List[dict[str, Any]] | List[KnowledgeBaseListEntry] | str | None:
+
+        if verbose and format:
+            logger.error("For knowledge base list, `--verbose` and `--format` are mutually exclusive options")
+            sys.exit(1)
+
         response = self.get_client().get()
         knowledge_bases = [KnowledgeBase.model_validate(knowledge_base) for knowledge_base in response]
 
@@ -234,7 +246,9 @@ class KnowledgeBaseController:
             for kb in knowledge_bases:
                 knowledge_base_list.append(json.loads(kb.model_dump_json(exclude_none=True)))
             rich.print(rich.json.JSON(json.dumps(knowledge_base_list, indent=4)))
+            return knowledge_base_list
         else:
+            knowledge_base_details=[]
             table = rich.table.Table(
                 show_header=True, 
                 header_style="bold white", 
@@ -251,6 +265,11 @@ class KnowledgeBaseController:
             for column in column_args:
                 table.add_column(column, **column_args[column])
             
+            connections_client = get_connections_client()
+            connections = connections_client.list()
+
+            connections_dict = {conn.connection_id: conn for conn in connections}
+            
             for kb in knowledge_bases:
                 app_id = ""
                 
@@ -258,18 +277,28 @@ class KnowledgeBaseController:
                    and kb.conversational_search_tool.index_config is not None \
                    and len(kb.conversational_search_tool.index_config) > 0 \
                    and kb.conversational_search_tool.index_config[0].connection_id is not None:
-                    connections_client = get_connections_client()
-                    app_id = str(connections_client.get_draft_by_id(kb.conversational_search_tool.index_config[0].connection_id))
-
-                table.add_row(
-                    kb.name,
-                    kb.description,
-                    app_id,
-                    str(kb.id)
+                    conn = connections_dict.get(kb.conversational_search_tool.index_config[0].connection_id)
+                    if conn:
+                        app_id = conn.app_id
+
+                entry = KnowledgeBaseListEntry(
+                    name=kb.name,
+                    id=str(kb.id),
+                    description=kb.description,
+                    app_id=app_id
                 )
-
-            rich.print(table)
-        
+                if format == ListFormats.JSON:
+                    knowledge_base_details.append(entry)
+                else:
+                    table.add_row(*entry.get_row_details())
+
+            match format:
+                case ListFormats.JSON:
+                    return knowledge_base_details
+                case ListFormats.Table:
+                    return rich_table_to_markdown(table)
+                case _:
+                    rich.print(table)   
 
     def remove_knowledge_base(self, id: str, name: str):
         knowledge_base_id = self.get_id(id, name)      

ibm_watsonx_orchestrate/cli/commands/models/models_command.py

@@ -136,11 +136,11 @@ def models_policy_import(
 def models_policy_add(
         name: Annotated[
             str,
-            typer.Option("--name", "-n", help="The name of the model to remove"),
+            typer.Option("--name", "-n", help="The name of the model policy to add"),
         ],
         models: Annotated[
             List[str],
-            typer.Option('--model', '-m', help='The name of the model to add'),
+            typer.Option('--model', '-m', help='A list of model names the policy should contain'),
         ],
         strategy: Annotated[
             ModelPolicyStrategyMode,