ibm-watsonx-orchestrate 1.6.3__py3-none-any.whl → 1.6.4__py3-none-any.whl

ibm_watsonx_orchestrate/cli/commands/copilot/copilot_controller.py

@@ -0,0 +1,368 @@
+import logging
+import os
+import sys
+import csv
+
+import rich
+from rich.console import Console
+from rich.prompt import Prompt
+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.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.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
+from ibm_watsonx_orchestrate.utils.exceptions import BadRequest
+
+logger = logging.getLogger(__name__)
+
+
+def _validate_output_file(output_file: str, dry_run_flag: bool) -> None:
+    if not output_file and not dry_run_flag:
+        logger.error(
+            "Please provide a valid yaml output file. Or use the `--dry-run` flag to output generated agent content to terminal")
+        sys.exit(1)
+
+    if output_file and dry_run_flag:
+        logger.error("Cannot set output file when performing a dry run")
+        sys.exit(1)
+
+    if output_file:
+        _, file_extension = os.path.splitext(output_file)
+        if file_extension not in {".yaml", ".yml", ".json"}:
+            logger.error("Output file must be of type '.yaml', '.yml' or '.json'")
+            sys.exit(1)
+
+
+def _get_progress_spinner() -> Progress:
+    console = Console()
+    return Progress(
+        SpinnerColumn(spinner_name="dots"),
+        TextColumn("[progress.description]{task.description}"),
+        transient=True,
+        console=console,
+    )
+
+
+def _get_incomplete_tool_from_name(tool_name: str) -> dict:
+    input_schema = ToolRequestBody(**{"type": "object", "properties": {}})
+    output_schema = ToolResponseBody(**{"description": "None"})
+    spec = ToolSpec(**{"name": tool_name, "description": tool_name, "permission": ToolPermission.ADMIN,
+                       "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_tools_from_names(tool_names: List[str]) -> List[dict]:
+    if not len(tool_names):
+        return []
+
+    tool_client = get_tool_client()
+
+    try:
+        with _get_progress_spinner() as progress:
+            task = progress.add_task(description="Fetching tools", total=None)
+            tools = tool_client.get_drafts_by_names(tool_names)
+            found_tools = {tool.get("name") for tool in tools}
+            progress.remove_task(task)
+            progress.refresh()
+            for tool_name in tool_names:
+                if tool_name not in found_tools:
+                    logger.warning(
+                        f"Failed to find tool named '{tool_name}'. Falling back to incomplete tool definition. Copilot performance maybe effected.")
+                    tools.append(_get_incomplete_tool_from_name(tool_name))
+    except ConnectionError:
+        logger.warning(
+            f"Failed to fetch tools from server. For optimal results please start the server and import the relevant tools {', '.join(tool_names)}.")
+        tools = []
+        for tool_name in tool_names:
+            tools.append(_get_incomplete_tool_from_name(tool_name))
+
+    return tools
+
+
+def _get_agents_from_names(collaborators_names: List[str]) -> List[dict]:
+    if not len(collaborators_names):
+        return []
+
+    native_agents_client = get_native_client()
+
+    try:
+        with _get_progress_spinner() as progress:
+            task = progress.add_task(description="Fetching agents", total=None)
+            agents = native_agents_client.get_drafts_by_names(collaborators_names)
+            found_agents = {tool.get("name") for tool in agents}
+            progress.remove_task(task)
+            progress.refresh()
+            for collaborator_name in collaborators_names:
+                if collaborator_name not in found_agents:
+                    logger.warning(
+                        f"Failed to find agent named '{collaborator_name}'. Falling back to incomplete agent definition. Copilot performance maybe effected.")
+                    agents.append(_get_incomplete_agent_from_name(collaborator_name))
+    except ConnectionError:
+        logger.warning(
+            f"Failed to fetch tools from server. For optimal results please start the server and import the relevant tools {', '.join(collaborators_names)}.")
+        agents = []
+        for collaborator_name in collaborators_names:
+            agents.append(_get_incomplete_agent_from_name(collaborator_name))
+
+    return agents
+
+def get_cpe_client() -> CPEClient:
+    url = os.getenv('CPE_URL', "http://localhost:8081")
+    return instantiate_client(client=CPEClient, url=url)
+
+
+def get_tool_client(*args, **kwargs):
+    return instantiate_client(ToolClient)
+
+
+def get_native_client(*args, **kwargs):
+    return instantiate_client(AgentClient)
+
+
+def gather_utterances(max: int) -> list[str]:
+    utterances = []
+    logger.info("Please provide 3 sample utterances you expect your agent to handle:")
+
+    count = 0
+
+    while count < max:
+        utterance = Prompt.ask("  [green]>[/green]").strip()
+
+        if utterance.lower() == 'q':
+            break
+
+        if utterance:
+            utterances.append(utterance)
+            count += 1
+
+    return utterances
+
+
+def get_deployed_tools_agents():
+    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}
+
+
+def pre_cpe_step(cpe_client):
+    tools_agents = get_deployed_tools_agents()
+    user_message = ""
+    with _get_progress_spinner() as progress:
+        task = progress.add_task(description="Initilizing Prompt Engine", total=None)
+        response = cpe_client.submit_pre_cpe_chat(user_message=user_message)
+        progress.remove_task(task)
+
+    res = {}
+    while True:
+        if "message" in response and response["message"]:
+            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"]:
+            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"]]
+            return res
+        with _get_progress_spinner() as progress:
+            task = progress.add_task(description="Thinking...", total=None)
+            response = cpe_client.submit_pre_cpe_chat(**message_content)
+            progress.remove_task(task)
+
+
+def find_tools_by_description(description, tool_client):
+    with _get_progress_spinner() as progress:
+        task = progress.add_task(description="Fetching Tools", total=None)
+        try:
+            tools = tool_client.get()
+            progress.remove_task(task)
+        except ConnectionError:
+            tools = []
+            progress.remove_task(task)
+            progress.refresh()
+            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)
+        try:
+            agents = agent_client.get()
+            progress.remove_task(task)
+        except ConnectionError:
+            agents = []
+            progress.remove_task(task)
+            progress.refresh()
+            logger.warning("Failed to contact wxo server to fetch agents. Proceeding with empty agent list")
+    return agents
+
+
+def gather_examples(samples_file=None):
+    if samples_file:
+        if samples_file.endswith('.txt'):
+            with open(samples_file) as f:
+                examples = f.read().split('\n')
+        elif samples_file.endswith('.csv'):
+            with open(samples_file, 'r', encoding='utf-8') as f:
+                reader = csv.DictReader(f)
+                if 'utterance' not in reader.fieldnames:
+                    raise BadRequest("CSV must have a column named 'utterance'")
+                examples = [row['utterance'].strip() for row in reader if row['utterance'].strip()]
+        else:
+            raise BadRequest(f'Unsupported samples file format: {os.path.basename(samples_file)}')
+    else:
+        examples = gather_utterances(3)
+
+    console = Console()
+    logger.info("You provided the following samples:")
+    for i, utterance in enumerate(examples, 1):
+        console.print(f"  {i}. {utterance}")
+
+    return examples
+
+
+def talk_to_cpe(cpe_client, samples_file=None, context_data=None):
+    context_data = context_data or {}
+    examples = gather_examples(samples_file)
+    # upload or gather input examples
+    context_data['examples'] = examples
+    response = None
+    with _get_progress_spinner() as progress:
+        task = progress.add_task(description="Thinking...", total=None)
+        response = cpe_client.init_with_context(context_data=context_data)
+        progress.remove_task(task)
+    accepted_prompt = None
+    while accepted_prompt is None:
+        resp = response.get('response')[0]
+        accepted_prompt = resp.get("final_zsh_prompt", None)
+        if not accepted_prompt:
+            cpe_message = resp.get("message", "")
+            rich.print('\n🤖 Copilot: ' + cpe_message)
+            message = Prompt.ask("\n👤 You").strip()
+            with _get_progress_spinner() as progress:
+                task = progress.add_task(description="Thinking...", total=None)
+                response = cpe_client.invoke(prompt=message)
+                progress.remove_task(task)
+
+    return accepted_prompt
+
+
+def prompt_tune(agent_spec: str, output_file: str | None, samples_file: str | None, dry_run_flag: bool) -> None:
+    agent = AgentsController.import_agent(file=agent_spec, app_id=None)[0]
+    agent_kind = agent.kind
+
+    if agent_kind != AgentKind.NATIVE:
+        logger.error(
+            f"Only native agents are supported for prompt tuning. Provided agent spec is on kind '{agent_kind}'")
+        sys.exit(1)
+
+    if not output_file and not dry_run_flag:
+        output_file = agent_spec
+
+    _validate_output_file(output_file, dry_run_flag)
+
+    client = get_cpe_client()
+
+    instr = agent.instructions
+
+    tools = _get_tools_from_names(agent.tools)
+
+    collaborators = _get_agents_from_names(agent.collaborators)
+    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})
+    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`")
+        sys.exit(1)
+
+    if new_prompt:
+        logger.info(f"The new instruction is: {new_prompt}")
+        agent.instructions = new_prompt
+
+        if dry_run_flag:
+            rich.print(agent.model_dump(exclude_none=True))
+        else:
+            if os.path.dirname(output_file):
+                os.makedirs(os.path.dirname(output_file), exist_ok=True)
+            AgentsController.persist_record(agent, output_file=output_file)
+
+
+def create_agent(output_file: str, llm: str, samples_file: str | None, dry_run_flag: bool = False) -> None:
+    _validate_output_file(output_file, dry_run_flag)
+    # 1. prepare the clients
+    cpe_client = get_cpe_client()
+
+    # 2. Pre-CPE stage:
+    try:
+        res = pre_cpe_step(cpe_client)
+    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`")
+        sys.exit(1)
+        
+    tools = res["tools"]
+    collaborators = res["collaborators"]
+    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})
+
+    # 6. create and save the agent
+    llm = llm if llm else DEFAULT_LLM
+    params = {
+        'style': agent_style,
+        'tools': [t['name'] for t in tools],
+        'llm': llm,
+        'collaborators': [c['name'] for c in collaborators]
+    }
+    agent = AgentsController.generate_agent_spec(agent_name, AgentKind.NATIVE, description, **params)
+    agent.instructions = instructions
+    agent.spec_version = SpecVersion.V1
+
+    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)
+    AgentsController.persist_record(agent, output_file=output_file)
+
+    message_lines = [
+        "Your agent building session finished successfully!",
+        f"Agent YAML saved in file:",
+        f"{os.path.abspath(output_file)}"
+    ]
+
+    # Determine the width of the frame
+    max_length = max(len(line) for line in message_lines)
+    frame_width = max_length + 4  # Padding for aesthetics
+
+    # Print the framed message
+    rich.print("╔" + "═" * frame_width + "╗")
+    for line in message_lines:
+        rich.print("║  " + line.ljust(max_length) + "  ║")
+    rich.print("╚" + "═" * frame_width + "╝")

ibm_watsonx_orchestrate/cli/commands/copilot/copilot_server_controller.py

@@ -0,0 +1,170 @@
+import logging
+import sys
+from pathlib import Path
+import subprocess
+import time
+import requests
+from urllib.parse import urlparse
+from ibm_watsonx_orchestrate.cli.commands.server.server_command import (
+    get_compose_file,
+    ensure_docker_compose_installed,
+    _prepare_clean_env,
+    ensure_docker_installed,
+    read_env_file,
+    get_default_env_file,
+    get_persisted_user_env,
+    get_dev_edition_source,
+    get_default_registry_env_vars_by_dev_edition_source,
+    docker_login_by_dev_edition_source,
+    write_merged_env_file,
+    apply_server_env_dict_defaults
+)
+
+logger = logging.getLogger(__name__)
+
+def wait_for_wxo_cpe_health_check(timeout_seconds=45, interval_seconds=2):
+    url = "http://localhost:8081/version"
+    logger.info("Waiting for Copilot component to be initialized...")
+    start_time = time.time()
+    while time.time() - start_time <= timeout_seconds:
+        try:
+            response = requests.get(url)
+            if 200 <= response.status_code < 300:
+                return True
+            else:
+                pass
+        except requests.RequestException as e:
+            pass
+
+        time.sleep(interval_seconds)
+    return False
+
+def _trim_authorization_urls(env_dict: dict) -> dict:
+    auth_url_key = "AUTHORIZATION_URL"
+    env_dict_copy = env_dict.copy()
+
+    auth_url = env_dict_copy.get(auth_url_key)
+    if not auth_url:
+        return env_dict_copy
+    
+    
+    parsed_url = urlparse(auth_url)
+    new_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+    env_dict_copy[auth_url_key] = new_url
+
+    return env_dict_copy
+
+    
+
+
+def run_compose_lite_cpe(user_env_file: Path) -> bool:
+    compose_path = get_compose_file()
+    compose_command = ensure_docker_compose_installed()
+    _prepare_clean_env(user_env_file)  
+    ensure_docker_installed()
+
+    default_env = read_env_file(get_default_env_file())
+    user_env = read_env_file(user_env_file) if user_env_file else {}
+    if not user_env:
+        user_env = get_persisted_user_env() or {}
+
+    dev_edition_source = get_dev_edition_source(user_env)
+    default_registry_vars = get_default_registry_env_vars_by_dev_edition_source(default_env, user_env, source=dev_edition_source)
+
+    # Update the default environment with the default registry variables only if they are not already set
+    for key in default_registry_vars:
+        if key not in default_env or not default_env[key]:
+            default_env[key] = default_registry_vars[key]
+
+    # Merge the default environment with the user environment
+    merged_env_dict = {
+        **default_env,
+        **user_env,
+    }
+    
+    merged_env_dict = apply_server_env_dict_defaults(merged_env_dict)
+    merged_env_dict = _trim_authorization_urls(merged_env_dict)
+
+    try:
+        docker_login_by_dev_edition_source(merged_env_dict, dev_edition_source)
+    except ValueError as ignored:
+        # do nothing, as the docker login here is not mandatory
+        pass
+
+    final_env_file = write_merged_env_file(merged_env_dict)
+
+    command = compose_command + [
+        "-f", str(compose_path),
+        "--env-file", str(final_env_file),
+        "up",
+        "cpe",
+        "-d",
+        "--remove-orphans"
+    ]
+
+    logger.info(f"Starting docker-compose Copilot service...")
+    result = subprocess.run(command, capture_output=False)
+
+    if result.returncode == 0:
+        logger.info("Copilot Service started successfully.")
+        # Remove the temp file if successful
+        if final_env_file.exists():
+            final_env_file.unlink()
+    else:
+        error_message = result.stderr.decode('utf-8') if result.stderr else "Error occurred."
+        logger.error(
+            f"Error running docker-compose (temporary env file left at {final_env_file}):\n{error_message}"
+        )
+        return False
+    
+    is_successful_cpe_healthcheck = wait_for_wxo_cpe_health_check()
+    if not is_successful_cpe_healthcheck:
+        logger.error("The Copilot service did not initialize within the expected time.  Check the logs for any errors.")
+
+    return True
+
+def run_compose_lite_cpe_down(is_reset: bool = False) -> None:
+    compose_path = get_compose_file()
+    compose_command = ensure_docker_compose_installed()
+    ensure_docker_installed()
+
+    default_env = read_env_file(get_default_env_file())
+    final_env_file = write_merged_env_file(default_env)
+
+    command = compose_command + [
+        "-f", str(compose_path),
+        "--env-file", final_env_file,
+        "down",
+        "cpe"
+    ]
+
+    if is_reset:
+        command.append("--volumes")
+        logger.info("Stopping docker-compose Copilot service and resetting volumes...")
+    else:
+        logger.info("Stopping docker-compose Copilot service...")
+
+    result = subprocess.run(command, capture_output=False)
+
+    if result.returncode == 0:
+        logger.info("Copilot service stopped successfully.")
+        # Remove the temp file if successful
+        if final_env_file.exists():
+            final_env_file.unlink()
+    else:
+        error_message = result.stderr.decode('utf-8') if result.stderr else "Error occurred."
+        logger.error(
+            f"Error running docker-compose (temporary env file left at {final_env_file}):\n{error_message}"
+        )
+        sys.exit(1)
+
+def start_server(user_env_file_path: Path) -> None:
+    is_server_started = run_compose_lite_cpe(user_env_file=user_env_file_path)
+
+    if is_server_started:
+        logger.info("Copilot service successfully started")
+    else:
+        logger.error("Unable to start orchestrate Copilot service.  Please check error messages and logs")
+
+def stop_server() -> None:
+    run_compose_lite_cpe_down()

ibm_watsonx_orchestrate/cli/commands/environment/environment_controller.py

@@ -23,7 +23,7 @@ from ibm_watsonx_orchestrate.cli.config import (
 )
 from ibm_watsonx_orchestrate.client.client import Client
 from ibm_watsonx_orchestrate.client.client_errors import ClientError
-from ibm_watsonx_orchestrate.client.agents.agent_client import AgentClient, ClientAPIException
+from ibm_watsonx_orchestrate.client.knowledge_bases.knowledge_base_client import KnowledgeBaseClient, ClientAPIException
 from ibm_watsonx_orchestrate.client.credentials import Credentials
 from threading import Lock
 from ibm_watsonx_orchestrate.client.utils import is_local_dev, check_token_validity, is_cpd_env
@@ -55,13 +55,13 @@ def _validate_token_functionality(token: str, url: str) -> None:
     '''
     is_cpd = is_cpd_env(url)
     if is_cpd is True:
-        agent_client = AgentClient(base_url=url, api_key=token, is_local=is_local_dev(url), verify=False)
+        knowledge_base_client = KnowledgeBaseClient(base_url=url, api_key=token, is_local=is_local_dev(url), verify=False)
     else:
-        agent_client = AgentClient(base_url=url, api_key=token, is_local=is_local_dev(url))
-    agent_client.api_key = token
+        knowledge_base_client = KnowledgeBaseClient(base_url=url, api_key=token, is_local=is_local_dev(url))
+    knowledge_base_client.api_key = token
 
     try:
-        agent_client.get()
+        knowledge_base_client.get()
     except ClientAPIException as e:
         if e.response.status_code >= 400:
             reason = e.response.reason

ibm_watsonx_orchestrate/cli/commands/environment/types.py

@@ -4,6 +4,8 @@ from enum import Enum
 class EnvironmentAuthType(str, Enum):
     IBM_CLOUD_IAM = 'ibm_iam'
     MCSP = 'mcsp'
+    MCSP_V1 = 'mcsp_v1'
+    MCSP_V2 = 'mcsp_v2'
     CPD = 'cpd'
 
     def __str__(self):