ibm-watsonx-orchestrate 1.12.2__py3-none-any.whl → 1.13.0b1__py3-none-any.whl

ibm_watsonx_orchestrate/cli/commands/evaluations/evaluations_environment_manager.py

@@ -0,0 +1,158 @@
+import logging
+import yaml
+from typing import Mapping, Any
+from enum import StrEnum
+from pathlib import Path
+
+from ibm_watsonx_orchestrate.cli.commands.agents.agents_controller import (
+    AgentsController,
+    Agent,
+    ExternalAgent,
+    AssistantAgent,
+)
+from ibm_watsonx_orchestrate.cli.commands.tools.tools_controller import (
+    ToolsController,
+    BaseTool,
+)
+from ibm_watsonx_orchestrate.cli.commands.knowledge_bases.knowledge_bases_controller import (
+    KnowledgeBaseController,
+    KnowledgeBase,
+)
+from ibm_watsonx_orchestrate.cli.commands.knowledge_bases.knowledge_bases_controller import (
+    parse_file as kb_parse_file,
+)
+from ibm_watsonx_orchestrate.cli.commands.evaluations.evaluations_controller import (
+    EvaluationsController,
+    EvaluateMode,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class ArtifactTypes(StrEnum):
+    """The allowed artifacts in the environment manager path.
+
+    The environment manager config looks like this:
+    ```json
+        env1:
+            agent:
+                agents_path: None
+            tools:
+                tools_path: None
+                tool_kind: None
+                # any other tool flags
+            knowledge:
+                knowledge_base_path: None
+        test_config: # path to config.yaml
+        clean_up: True
+    ```
+    The allowed artifacts/keys are "agent", "tools", "knowledge"
+    """
+
+    agent = "agent"
+    tools = "tools"
+    knowledge = "knowledge"
+
+
+class TestCaseManager:
+    def __init__(
+        self,
+        env_settings: Mapping[str, Any],
+        output_dir: str,
+        mode: EvaluateMode = EvaluateMode.default,
+    ):
+        self.env_settings = env_settings
+        self.cleanup = env_settings.get("clean_up", False)
+        self.output_dir = output_dir
+        self.mode = mode
+
+        self.agent_controller = AgentsController()
+        self.knowledge_controller = KnowledgeBaseController()
+        self.tool_controller = None
+        if (tool_settings := env_settings.get(ArtifactTypes.tools)):
+            self.tool_controller = ToolsController(
+                tool_kind=tool_settings.get("kind"),
+                file=tool_settings.get("file"),
+                requirements_file=tool_settings.get("requirements_file")
+            )
+
+        self.imported_artifacts = []
+
+    def __enter__(self):
+        for artifact in [
+            ArtifactTypes.tools,
+            ArtifactTypes.knowledge,
+            ArtifactTypes.agent,
+        ]:
+            if artifact not in self.env_settings:
+                continue
+
+            artifact_settings = self.env_settings.get(artifact)
+            if artifact == ArtifactTypes.tools:
+                tools = ToolsController.import_tool(**artifact_settings)
+                # import_tool returns Iterator[BaseTool], copy the iterator into a list for preservation
+                # this is needed if user wants environment cleanup
+                tools = [tool for tool in tools]
+                self.imported_artifacts.append(tools)
+                self.tool_controller.publish_or_update_tools(tools)
+            elif artifact == ArtifactTypes.knowledge:
+                KnowledgeBaseController.import_knowledge_base(**artifact_settings)
+                kb_spec = kb_parse_file(artifact_settings.get("file"))
+                self.imported_artifacts.append(kb_spec)
+            elif artifact == ArtifactTypes.agent:
+                artifact_settings["app_id"] = artifact_settings.get("app_id", None)
+                agents = AgentsController.import_agent(**artifact_settings)
+                self.agent_controller.publish_or_update_agents(agents)
+                self.imported_artifacts.append(agents)
+
+        eval = EvaluationsController()
+        eval.evaluate(
+            test_paths=self.env_settings.get("test_paths"),
+            output_dir=self.output_dir,
+            mode=self.mode,
+        )
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.cleanup:
+            logger.info("Cleaning environment")
+            for artifact in self.imported_artifacts:
+                # artifact can be a list of agents, tools
+                for item in artifact:
+                    if isinstance(item, BaseTool):
+                        self.tool_controller.remove_tool(item.__tool_spec__.name)
+                    if isinstance(item, KnowledgeBase):
+                        self.knowledge_controller.remove_knowledge_base(
+                            item.id, item.name
+                        )
+                    if isinstance(item, (Agent, AssistantAgent, ExternalAgent)):
+                        self.agent_controller.remove_agent(item.name, item.kind)
+
+
+def run_environment_manager(
+    environment_manager_path: str,
+    mode: EvaluateMode = EvaluateMode.default,
+    output_dir: str = None,
+):
+    with open(environment_manager_path, encoding="utf-8", mode="r") as f:
+        env_settings = yaml.load(f, Loader=yaml.SafeLoader)
+
+    for env in env_settings:
+        if not env_settings.get(env).get("enabled"):
+            continue
+        results_folder = Path(output_dir) / env
+        results_folder.mkdir(parents=True, exist_ok=True)
+        logger.info(
+            "Processing environment: '%s'. Results will be saved to '%s'",
+            env,
+            results_folder,
+        )
+
+        with TestCaseManager(
+            env_settings=env_settings.get(env),
+            output_dir=str(results_folder),
+            mode=mode,
+        ):
+            logger.info("Finished evaluation for environment: '%s'", env)

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_command.py

@@ -59,3 +59,29 @@ def knowledge_base_status(
 ):  
     controller = KnowledgeBaseController()
     controller.knowledge_base_status(id=id, name=name)
+
+@knowledge_bases_app.command(name="export", help='Export a knowledge base spec to a yaml')
+def knowledge_base_export(
+    output_file: Annotated[
+        str,
+        typer.Option(
+            "--output",
+            "-o",
+            help="Path to a where the zip file containing the exported data should be saved",
+        ),
+    ],
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="The name of the knowledge base you want to export"),
+    ]=None,
+    id: Annotated[
+        str,
+        typer.Option("--id", "-i", help="The ID of the knowledge base you wish export"),
+    ]=None,
+):
+    controller = KnowledgeBaseController()
+    controller.knowledge_base_export(
+        id=id,
+        name=name,
+        output_path=output_file
+    )

ibm_watsonx_orchestrate/cli/commands/knowledge_bases/knowledge_bases_controller.py

@@ -5,15 +5,20 @@ import requests
 import logging
 import importlib
 import inspect
+import yaml
 from pathlib import Path
-from typing import List
+from typing import List, Any, Optional
+from zipfile import ZipFile
+from io import BytesIO
 
 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
+from ibm_watsonx_orchestrate.agent_builder.knowledge_bases.types import KnowledgeBaseKind, IndexConnection, SpecVersion
 
 logger = logging.getLogger(__name__)
 
@@ -63,6 +68,32 @@ def build_file_object(file_dir: str, file: str | FileUpload):
         return ('files', (get_file_name(file.path), open(get_relative_file_path(file.path, file_dir), 'rb')))
     return ('files', (get_file_name(file), open(get_relative_file_path(file, file_dir), 'rb')))
 
+def build_connections_map(key_attr: str) -> dict:
+    connections_client = get_connections_client()
+    connections = connections_client.list()
+
+    return {getattr(conn, key_attr): conn for conn in connections}
+
+def get_index_config(kb: KnowledgeBase, index: int = 0) -> IndexConnection | None:
+    if kb.conversational_search_tool is not None \
+        and kb.conversational_search_tool.index_config is not None \
+        and len(kb.conversational_search_tool.index_config) > index:
+
+        return kb.conversational_search_tool.index_config[index]
+    return None
+
+def get_kb_app_id(kb: KnowledgeBase) -> str | None:
+    index_config = get_index_config(kb)
+    if not index_config:
+        return
+    return index_config.app_id
+
+def get_kb_connection_id(kb: KnowledgeBase) -> str | None:
+    index_config = get_index_config(kb)
+    if not index_config:
+        return
+    return index_config.connection_id
+
 class KnowledgeBaseController:
     def __init__(self):
         self.client = None
@@ -78,24 +109,23 @@ class KnowledgeBaseController:
 
         knowledge_bases = parse_file(file=file)
         
-        if app_id:
-            connections_client = get_connections_client()
-            connection_id = None
-                
-            connections = connections_client.get_draft_by_app_id(app_id=app_id)
-            if not connections:
-                logger.error(f"No connection exists with the app-id '{app_id}'")
-                exit(1)
-
-            connection_id = connections.connection_id
-
-            for kb in knowledge_bases:
-                if kb.conversational_search_tool and kb.conversational_search_tool.index_config and len(kb.conversational_search_tool.index_config) > 0:
-                    kb.conversational_search_tool.index_config[0].connection_id = connection_id
+        connections_map = None
         
         existing_knowledge_bases = client.get_by_names([kb.name for kb in knowledge_bases])
         
         for kb in knowledge_bases:
+            app_id = app_id if app_id else get_kb_app_id(kb)
+            if app_id:
+                if not connections_map:
+                    connections_map = build_connections_map("app_id")
+                conn = connections_map.get(app_id)
+                if conn:
+                    index_config = get_index_config(kb)
+                    if index_config:
+                        index_config.connection_id = conn.connection_id
+                else:
+                    logger.error(f"No connection exists with the app-id '{app_id}'")
+                    exit(1)
             try:
                 file_dir = "/".join(file.split("/")[:-1])
 
@@ -198,8 +228,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 +248,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 +275,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,25 +294,34 @@ class KnowledgeBaseController:
             for column in column_args:
                 table.add_column(column, **column_args[column])
             
+            connections_dict = build_connections_map("connection_id")
+            
             for kb in knowledge_bases:
                 app_id = ""
-                
-                if kb.conversational_search_tool is not None \
-                   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)
+                connection_id = get_kb_connection_id(kb)
+                if connection_id is not None:
+                    conn = connections_dict.get(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
                 )
+                if format == ListFormats.JSON:
+                    knowledge_base_details.append(entry)
+                else:
+                    table.add_row(*entry.get_row_details())
 
-            rich.print(table)
-        
+            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)      
@@ -283,4 +335,68 @@ class KnowledgeBaseController:
                 logger.warning(f"No knowledge base {logEnding} found")
             logger.error(e.response.text)
             exit(1)
+    
+    def get_knowledge_base(self, id) -> KnowledgeBase:
+        client = self.get_client()
+        try:
+            return KnowledgeBase.model_validate(client.get_by_id(id))
+        except requests.HTTPError as e:
+            if e.response.status_code == 404:
+                logger.error(f"No knowledge base {id} found")
+            else:
+                logger.error(e.response.text)
+            exit(1)
+
+
+    def knowledge_base_export(self,
+            output_path: str,
+            id: Optional[str] = None,
+            name: Optional[str] = None, 
+            zip_file_out: Optional[ZipFile] = None) -> None:
+        output_file = Path(output_path)
+        output_file_extension = output_file.suffix
+        if output_file_extension not in  {".yaml", ".yml"} :
+            logger.error(f"Output file must end with the extension '.yaml'/'.yml'. Provided file '{output_path}' ends with '{output_file_extension}'")
+            sys.exit(1)
+        
+        knowledge_base_id = self.get_id(id, name)
+        logEnding = f"with ID '{id}'" if id else f"'{name}'"  
+        
+        logger.info(f"Exporting spec for knowledge base {logEnding}'")
+
+        knowledge_base = self.get_knowledge_base(knowledge_base_id)
 
+        if not knowledge_base:
+            logger.error(f"Knowledge base'{knowledge_base_id}' not found.'")
+            return
+        
+        knowledge_base.tenant_id = None
+        knowledge_base.id = None
+        knowledge_base.spec_version = SpecVersion.V1
+        knowledge_base.kind = KnowledgeBaseKind.KNOWLEDGE_BASE
+        
+        connection_id = get_kb_connection_id(knowledge_base)
+        if connection_id:
+            connections_map = build_connections_map("connection_id")
+            conn = connections_map.get(connection_id)
+            if conn:
+                index_config = get_index_config(knowledge_base)
+                index_config.app_id = conn.app_id
+                index_config.connection_id = None
+            else:
+                logger.warning(f"Connection '{connection_id}' not found, unable to resolve app_id for Knowledge base {logEnding}")
+
+        knowledge_base_spec = knowledge_base.model_dump(mode="json", exclude_none=True, exclude_unset=True)
+        if zip_file_out:
+            knowledge_base_spec_yaml = yaml.dump(knowledge_base_spec, sort_keys=False, default_flow_style=False, allow_unicode=True)
+            knowledge_base_spec_yaml_bytes = knowledge_base_spec_yaml.encode("utf-8")
+            knowledge_base_spec_yaml_file = BytesIO(knowledge_base_spec_yaml_bytes)
+            zip_file_out.writestr(
+                output_path,
+                knowledge_base_spec_yaml_file.getvalue()
+            )
+        else:
+            with open(output_path, 'w') as outfile:
+                yaml.dump(knowledge_base_spec, outfile, sort_keys=False, default_flow_style=False, allow_unicode=True)
+        
+        logger.info(f"Successfully exported for knowledge base {logEnding} to '{output_path}'")

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,

ibm_watsonx_orchestrate/cli/commands/models/models_controller.py

@@ -6,7 +6,7 @@ import yaml
 import importlib
 import inspect
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
 import requests
 import rich
@@ -16,10 +16,11 @@ from ibm_watsonx_orchestrate.client.model_policies.model_policies_client import
 from ibm_watsonx_orchestrate.agent_builder.model_policies.types import ModelPolicy, ModelPolicyInner, \
     ModelPolicyRetry, ModelPolicyStrategy, ModelPolicyStrategyMode, ModelPolicyTarget
 from ibm_watsonx_orchestrate.client.models.models_client import ModelsClient
-from ibm_watsonx_orchestrate.agent_builder.models.types import VirtualModel, ProviderConfig, ModelType, ANTHROPIC_DEFAULT_MAX_TOKENS
+from ibm_watsonx_orchestrate.agent_builder.models.types import VirtualModel, ProviderConfig, ModelType, ANTHROPIC_DEFAULT_MAX_TOKENS, ModelListEntry
 from ibm_watsonx_orchestrate.client.utils import instantiate_client, is_cpd_env
 from ibm_watsonx_orchestrate.client.connections import get_connection_id, ConnectionType
 from ibm_watsonx_orchestrate.utils.environment import EnvService
+from ibm_watsonx_orchestrate.cli.common import ListFormats, rich_table_to_markdown
 
 logger = logging.getLogger(__name__)
 
@@ -149,7 +150,7 @@ class ModelsController:
             self.model_policies_client = instantiate_client(ModelPoliciesClient)
         return self.model_policies_client
 
-    def list_models(self, print_raw: bool = False) -> None:
+    def list_models(self, print_raw: bool = False, format: Optional[ListFormats] = None) -> List[ModelListEntry] | str |None:
         models_client: ModelsClient = self.get_models_client()
         model_policies_client: ModelPoliciesClient = self.get_model_policies_client()
         global WATSONX_URL
@@ -224,6 +225,7 @@ class ModelsController:
 
                 console.print("[yellow]★[/yellow] [italic dim]indicates a supported and preferred model[/italic dim]\n[blue dim]✨️[/blue dim] [italic dim]indicates a model from a custom provider[/italic dim]" )
             else:
+                model_details = []
                 table = rich.table.Table(
                     show_header=True,
                     title="[bold]Available Models[/bold]",
@@ -234,15 +236,32 @@ class ModelsController:
                     table.add_column(col)
 
                 for model in (virtual_models + virtual_model_policies):
-                    table.add_row(f"✨️ {model.name}", model.description or 'No description provided.')
+                    entry = ModelListEntry(
+                        name=model.name,
+                        description=model.description,
+                        is_custom=True
+                    )
+                    model_details.append(entry)
+                    table.add_row(*entry.get_row_details())
 
                 for model in sorted_models:
-                    model_id = model.get("model_id", "N/A")
-                    short_desc = model.get("short_description", "No description provided.")
-                    marker = "★ " if any(pref in model_id.lower() for pref in preferred_list) else ""
-                    table.add_row(f"[yellow]{marker}[/yellow]watsonx/{model_id}", short_desc)
-            
-                rich.print(table)
+                    name = model.get("model_id", "N/A")
+                    entry = ModelListEntry(
+                        name=name,
+                        description=model.get("short_description"),
+                        is_custom=False,
+                        recommended=any(pref in name.lower() for pref in preferred_list)
+                    )
+                    model_details.append(entry)
+                    table.add_row(*entry.get_row_details())
+
+                match format:
+                    case ListFormats.JSON:
+                        return model_details
+                    case ListFormats.Table:
+                        return rich_table_to_markdown(table)
+                    case _: 
+                        rich.print(table)
 
     def import_model(self, file: str, app_id: str | None) -> List[VirtualModel]:
         from ibm_watsonx_orchestrate.cli.commands.models.model_provider_mapper import validate_ProviderConfig # lazily import this because the lut building is expensive