ibm-watsonx-orchestrate 1.2.0__py3-none-any.whl → 1.4.2__py3-none-any.whl

ibm_watsonx_orchestrate/cli/commands/toolkit/toolkit_controller.py

@@ -8,6 +8,9 @@ import sys
 import re
 import requests
 from ibm_watsonx_orchestrate.client.toolkit.toolkit_client import ToolKitClient
+from ibm_watsonx_orchestrate.client.tools.tool_client import ToolClient
+from ibm_watsonx_orchestrate.agent_builder.toolkits.base_toolkit import BaseToolkit, ToolkitSpec
+from ibm_watsonx_orchestrate.agent_builder.toolkits.types import ToolkitKind, Language, ToolkitSource
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
 from ibm_watsonx_orchestrate.utils.utils import sanatize_app_id
 from ibm_watsonx_orchestrate.client.connections import get_connections_client
@@ -16,12 +19,13 @@ import json
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from ibm_watsonx_orchestrate.client.utils import is_local_dev
+from rich.json import JSON
+import rich
+import rich.table
+import json
 
 logger = logging.getLogger(__name__)
 
-class ToolkitKind(str, Enum):
-    MCP = "mcp"
-
 def get_connection_id(app_id: str) -> str:
     connections_client = get_connections_client()
     existing_draft_configuration = connections_client.get_config(app_id=app_id, env='draft')
@@ -47,29 +51,34 @@ def validate_params(kind: str):
 
 class ToolkitController:
     def __init__(
-        self, 
+        self,
         kind: ToolkitKind = None,
         name: str = None,
         description: str = None,
-        package_root: str  = None,
-        command: str  = None,
+        package: str = None,
+        package_root: str = None,
+        language: Language = None,
+        command: str = None,
     ):
         self.kind = kind
         self.name = name
         self.description = description
+        self.package = package
         self.package_root = package_root
+        self.language = language
         self.command = command
         self.client = None
 
+        self.source: ToolkitSource = (
+            ToolkitSource.PUBLIC_REGISTRY if package else ToolkitSource.FILES
+        )
+
     def get_client(self) -> ToolKitClient:
         if not self.client:
             self.client = instantiate_client(ToolKitClient)
         return self.client
 
     def import_toolkit(self, tools: Optional[List[str]] = None, app_id: Optional[List[str]] = None):
-        if not is_local_dev():
-            logger.error("This functionality is only available for Local Environments")
-            sys.exit(1)
 
         if app_id and isinstance(app_id, str):
             app_id = [app_id]
@@ -86,26 +95,28 @@ class ToolkitController:
             logger.error(f"Existing toolkit found with name '{self.name}'. Failed to create toolkit.")
             sys.exit(1)
 
-        with tempfile.TemporaryDirectory() as tmpdir:
-            # Handle zip file or directory
-            if self.package_root.endswith(".zip") and os.path.isfile(self.package_root):
-                zip_file_path = self.package_root
-            else:
-                zip_file_path = os.path.join(tmpdir, os.path.basename(f"{self.package_root.rstrip(os.sep)}.zip"))
-                with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as mcp_zip_tool_artifacts:
-                    self._populate_zip(self.package_root, mcp_zip_tool_artifacts)
+        try:
+            command_parts = json.loads(self.command)
+            if not isinstance(command_parts, list):
+                raise ValueError("JSON command must be a list of strings")
+        except (json.JSONDecodeError, ValueError):
+            command_parts = self.command.split()
+
+        command = command_parts[0]
+        args = command_parts[1:]
 
-            try:
-                command_parts = json.loads(self.command)
-                if not isinstance(command_parts, list):
-                    raise ValueError("JSON command must be a list of strings")
-            except (json.JSONDecodeError, ValueError):
-                command_parts = self.command.split()
+        console = Console()
 
-            command = command_parts[0]
-            args = command_parts[1:]
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Handle zip file or directory
+            if self.package_root:
+                if self.package_root.endswith(".zip") and os.path.isfile(self.package_root):
+                    zip_file_path = self.package_root
+                else:
+                    zip_file_path = os.path.join(tmpdir, os.path.basename(f"{self.package_root.rstrip(os.sep)}.zip"))
+                    with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as mcp_zip_tool_artifacts:
+                        self._populate_zip(self.package_root, mcp_zip_tool_artifacts)
 
-            console = Console()
             # List tools if not provided
             if tools is None:
                 with Progress(
@@ -116,48 +127,59 @@ class ToolkitController:
                 ) as progress:
                     progress.add_task(description="No tools specified, retrieving all tools from provided MCP server", total=None)
                     tools = self.get_client().list_tools(
-                    zip_file_path=zip_file_path,
-                    command=command,
-                    args=args,
-                )
-                # Normalize tools to a list of tool names
+                        zip_file_path=zip_file_path,
+                        command=command,
+                        args=args,
+                    )
+                
                 tools = [
                     tool["name"] if isinstance(tool, dict) and "name" in tool else tool
                     for tool in tools
                 ]
 
-
                 logger.info("✅ The following tools will be imported:")
                 for tool in tools:
                     console.print(f"  • {tool}")
 
-
             # Create toolkit metadata
             payload = {
                 "name": self.name,
                 "description": self.description,
                 "mcp": {
-                    "source": "files",
+                    "source": self.source.value,
                     "command": command,
                     "args": args,
                     "tools": tools,
                     "connections": remapped_connections,
                 }
             }
-            toolkit = self.get_client().create_toolkit(payload)
-            toolkit_id = toolkit["id"]
 
-            console = Console()
-            # Upload zip file
+
             with Progress(
                 SpinnerColumn(spinner_name="dots"),
                 TextColumn("[progress.description]{task.description}"),
                 transient=True,
                 console=console,
             ) as progress:
-                progress.add_task(description="Uploading toolkit zip file...", total=None)
-                self.get_client().upload(toolkit_id=toolkit_id, zip_file_path=zip_file_path)
-            logger.info(f"Successfully imported tool kit {self.name}")
+                progress.add_task(description="Creating toolkit...", total=None)
+                toolkit = self.get_client().create_toolkit(payload)
+
+            toolkit_id = toolkit["id"]
+
+
+
+            # Upload zip file
+            if self.package_root:
+                with Progress(
+                    SpinnerColumn(spinner_name="dots"),
+                    TextColumn("[progress.description]{task.description}"),
+                    transient=True,
+                    console=console,
+                ) as progress:
+                    progress.add_task(description="Uploading toolkit zip file...", total=None)
+                    self.get_client().upload(toolkit_id=toolkit_id, zip_file_path=zip_file_path)
+
+        logger.info(f"Successfully imported tool kit {self.name}")
 
     def _populate_zip(self, package_root: str, zipfile: zipfile.ZipFile) -> str:
         for root, _, files in os.walk(package_root):
@@ -209,4 +231,63 @@ class ToolkitController:
                 logger.warning(f"No toolkit named '{name}' found")
         except requests.HTTPError as e:
             logger.error(e.response.text)
-            exit(1)
+            exit(1)
+
+    def list_toolkits(self, verbose=False):
+        client = self.get_client()
+        response = client.get()
+        toolkit_spec = [ToolkitSpec.model_validate(toolkit) for toolkit in response]
+        toolkits = [BaseToolkit(spec=spec) for spec in toolkit_spec]
+
+        if verbose:
+            tools_list = []
+            for toolkit in toolkits:
+                tools_list.append(json.loads(toolkit.dumps_spec()))
+            rich.print(JSON(json.dumps(tools_list, indent=4)))
+        else:
+            table = rich.table.Table(show_header=True, header_style="bold white", show_lines=True)
+            columns = ["Name", "Kind", "Description", "Tools", "App ID"]
+            for column in columns:
+                table.add_column(column)
+
+            tools_client = instantiate_client(ToolClient)
+
+            connections_client = get_connections_client()
+            connections = connections_client.list()
+
+            connections_dict = {conn.connection_id: conn for conn in connections}
+
+            for toolkit in toolkits:
+                tool_ids = toolkit.__toolkit_spec__.tools or []
+                tool_names = []
+                if len(tool_ids) == 0:
+                    logger.warning("This toolkit contains no tools.")
+
+                for tool_id in tool_ids:
+                    tool = tools_client.get_draft_by_id(tool_id)
+                    tool_names.append(tool["name"])
+
+                app_ids = []
+                connection_ids = toolkit.__toolkit_spec__.mcp.connections.values()
+
+                for connection_id in connection_ids:
+                    connection = connections_dict.get(connection_id)
+                    if connection:
+                        app_id = str(connection.app_id or connection.connection_id)
+                    elif connection_id:
+                        app_id = str(connection_id)
+                    else:
+                        app_id = ""
+                    app_ids.append(app_id)
+
+
+
+                table.add_row(
+                    toolkit.__toolkit_spec__.name,
+                    "MCP",
+                    toolkit.__toolkit_spec__.description,
+                    ", ".join(tool_names),
+                    ", ".join(app_ids),
+                )
+
+            rich.print(table)

ibm_watsonx_orchestrate/cli/commands/tools/tools_command.py

@@ -15,7 +15,7 @@ def tool_import(
         typer.Option(
             "--file",
             "-f",
-            help="Path to Python or OpenAPI spec YAML file. Required for kind openapi or python",
+            help="Path to Python, OpenAPI spec YAML file or flow JSON or python file. Required for kind openapi, python and flow",
         ),
     ] = None,
     # skillset_id: Annotated[
@@ -83,3 +83,24 @@ def remove_tool(
 ):  
     tools_controller = ToolsController()
     tools_controller.remove_tool(name=name)
+
+@tools_app.command(name="export", help='Export a tool to a zip file')
+def tool_export(
+    name: Annotated[
+        str,
+        typer.Option("--name", "-n", help="The name of the tool you want to export"),
+    ],
+    output_file: Annotated[
+        str,
+        typer.Option(
+            "--output",
+            "-o",
+            help="Path to a where the zip file containing the exported data should be saved",
+        ),
+    ],
+):
+    tools_controller = ToolsController()
+    tools_controller.export_tool(
+        name=name,
+        output_path=output_file
+    )

ibm_watsonx_orchestrate/cli/commands/tools/tools_controller.py

@@ -3,6 +3,7 @@ import asyncio
 import importlib
 import inspect
 import sys
+import io
 import re
 import tempfile
 import requests
@@ -19,18 +20,28 @@ import glob
 import rich.table
 import typer
 
+from rich.console import Console
+from rich.panel import Panel
+
 from ibm_watsonx_orchestrate.agent_builder.tools import BaseTool, ToolSpec
-from ibm_watsonx_orchestrate.agent_builder.tools.openapi_tool import create_openapi_json_tools_from_uri
+from ibm_watsonx_orchestrate.agent_builder.tools.openapi_tool import create_openapi_json_tools_from_uri,create_openapi_json_tools_from_content
+from ibm_watsonx_orchestrate.cli.commands.models.models_command import ModelHighlighter
 from ibm_watsonx_orchestrate.cli.commands.tools.types import RegistryType
+from ibm_watsonx_orchestrate.cli.commands.connections.connections_controller import configure_connection, remove_connection, add_connection
+from ibm_watsonx_orchestrate.agent_builder.connections.types import  ConnectionType, ConnectionEnvironment, ConnectionPreference
 from ibm_watsonx_orchestrate.cli.config import Config, CONTEXT_SECTION_HEADER, CONTEXT_ACTIVE_ENV_OPT, \
     PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TYPE_OPT, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT, \
     DEFAULT_CONFIG_FILE_CONTENT
 from ibm_watsonx_orchestrate.agent_builder.connections import ConnectionSecurityScheme, ExpectedCredentials
+from ibm_watsonx_orchestrate.experimental.flow_builder.flows.decorators import FlowWrapper
 from ibm_watsonx_orchestrate.client.tools.tool_client import ToolClient
 from ibm_watsonx_orchestrate.client.toolkit.toolkit_client import ToolKitClient
 from ibm_watsonx_orchestrate.client.connections import get_connections_client, get_connection_type
 from ibm_watsonx_orchestrate.client.utils import instantiate_client, is_local_dev
 from ibm_watsonx_orchestrate.utils.utils import sanatize_app_id
+from ibm_watsonx_orchestrate.client.utils import is_local_dev
+from ibm_watsonx_orchestrate.client.tools.tempus_client import TempusClient
+from ibm_watsonx_orchestrate.experimental.flow_builder.utils import import_flow_model
 
 from  ibm_watsonx_orchestrate import __version__
 
@@ -43,19 +54,20 @@ class ToolKind(str, Enum):
     openapi = "openapi"
     python = "python"
     mcp = "mcp"
+    flow = "flow"
     # skill = "skill"
 
 def validate_app_ids(kind: ToolKind, **args) -> None:
     app_ids = args.get("app_id")
     if not app_ids:
         return
-    
+
     if kind == ToolKind.openapi:
         if app_ids and len(app_ids) > 1:
             raise typer.BadParameter(
                 "Kind 'openapi' can only take one app-id"
             )
-    
+
     connections_client = get_connections_client()
 
     imported_connections_list = connections_client.list()
@@ -115,7 +127,7 @@ def get_connection_id(app_id: str) -> str:
 
 def parse_python_app_ids(app_ids: List[str]) -> dict[str,str]:
     app_id_dict = {}
-    for app_id in app_ids:        
+    for app_id in app_ids:
         # Split on = but not on \=
         split_pattern = re.compile(r"(?<!\\)=")
         split_id = re.split(split_pattern, app_id)
@@ -139,7 +151,7 @@ def parse_python_app_ids(app_ids: List[str]) -> dict[str,str]:
 def validate_python_connections(tool: BaseTool):
     if not tool.expected_credentials:
         return
-    
+
     connections_client = get_connections_client()
     connections = tool.__tool_spec__.binding.python.connections
 
@@ -158,33 +170,32 @@ def validate_python_connections(tool: BaseTool):
             expected_tool_conn_types = expected_cred.type
         else:
             expected_tool_conn_types = [expected_cred.type]
-        
+
         sanatized_expected_tool_app_id = sanatize_app_id(expected_tool_app_id)
         if sanatized_expected_tool_app_id in existing_sanatized_expected_tool_app_ids:
             logger.error(f"Duplicate App ID found '{expected_tool_app_id}'. Multiple expected app ids in the tool '{tool.__tool_spec__.name}' collide after sanaitization to '{sanatized_expected_tool_app_id}'. Please rename the offending app id in your tool.")
             sys.exit(1)
         existing_sanatized_expected_tool_app_ids.add(sanatized_expected_tool_app_id)
-        
+
         if sanatized_expected_tool_app_id not in provided_connections:
             logger.error(f"The tool '{tool.__tool_spec__.name}' requires an app-id '{expected_tool_app_id}'. Please use the `--app-id` flag to provide the required app-id")
-            validation_failed = True
+            sys.exit(1)
 
         if not connections:
             continue
-
+            
         connection_id = connections.get(sanatized_expected_tool_app_id)
-        
         imported_connection = imported_connections.get(connection_id)
         imported_connection_auth_type = get_connection_type(security_scheme=imported_connection.security_scheme, auth_type=imported_connection.auth_type)
 
         if connection_id and not imported_connection:
-            logger.error(f"The expected connection id '{connection_id}' does not match any known connection. This is likely caused by the connection being delted. Please rec-reate the connection and re-import the tool")
+            logger.error(f"The expected connection id '{connection_id}' does not match any known connection. This is likely caused by the connection being deleted. Please rec-reate the connection and re-import the tool")
             validation_failed = True
 
         if imported_connection and len(expected_tool_conn_types) and imported_connection_auth_type not in expected_tool_conn_types:
             logger.error(f"The app-id '{imported_connection.app_id}' is of type '{imported_connection_auth_type.value}'. The tool '{tool.__tool_spec__.name}' accepts connections of the following types '{', '.join(expected_tool_conn_types)}'. Use `orchestrate connections list` to view current connections and use `orchestrate connections add` to create the relevent connection")
             validation_failed = True
-        
+
     if validation_failed:
         exit(1)
 
@@ -302,22 +313,141 @@ def import_python_tool(file: str, requirements_file: str = None, app_id: List[st
             obj.__tool_spec__.binding.python.function = f"{file_name.replace('.py', '')}:{fn}"
 
         else:
-            package = package[1:]
+            pkg = package[1:]
             fn = obj.__tool_spec__.binding.python.function[obj.__tool_spec__.binding.python.function.index(':')+1:]
-            obj.__tool_spec__.binding.python.function = f"{package}:{fn}"
+            obj.__tool_spec__.binding.python.function = f"{pkg}:{fn}"
 
         if app_id and len(app_id):
             obj.__tool_spec__.binding.python.connections = parse_python_app_ids(app_id)
 
         validate_python_connections(obj)
         tools.append(obj)
-    
+
     return tools
 
+async def import_flow_tool(file: str) -> None:
+    
+    '''
+    Import a flow tool from a file. The file can be either a python file or a json file.
+    If the file is a python file, it should contain a flow model builder function decorated with the @flow decorator.
+    If the file is a json file, it should contain a flow model in json format.
+    Also, a connection will be created for the flow if one does not exists and the environment token will be used.  This is a 
+    workaround until flow bindings are supported in the server.
+    The function will return a list of tools created from the flow model.
+    '''
+
+    theme = rich.theme.Theme({"model.name": "bold cyan"})
+    console = rich.console.Console(highlighter=ModelHighlighter(), theme=theme)
+    
+    message = f"""[bold cyan]Flow Tools: Experimental Feature[/bold cyan]
+   
+The [bold]flow tool[/bold] is being imported from [green]`{file}`[/green].  
+    
+[bold cyan]Additional information:[/bold cyan]
+
+- Ensure the flow engine is running by issuing the [bold cyan]orchestrate server start[/bold cyan] command with the [bold cyan]--with-flow-runtime[/bold cyan] option
+- The [bold green]get_flow_status[/bold green] tool is being imported to support flow tools. Ensure [bold]both this tools and the one you are importing are added to your agent[/bold] to retrieve the flow output. 
+- Include additional instructions in your agent to call the [bold green]get_flow_status[/bold green] tool to retrieve the flow output. For example: [green]"If you get an instance_id, use the tool get_flow_status to retrieve the current status of a flow."[/green]
+
+    """
+
+    console.print(Panel(message,  title="[bold blue]Flow tool support information[/bold blue]", border_style="bright_blue"))
+   
+
+    if not is_local_dev():
+        raise typer.BadParameter(f"Flow tools are only supported in local environment.")
+
+    model = None
+    
+    # Load the Flow JSON model from the file
+    try:
+        file_path = Path(file).absolute()
+        file_path_str = str(file_path)
+
+        if file_path.is_dir():
+            raise typer.BadParameter(f"Provided flow file path is not a file.")
+
+        elif file_path.is_symlink():
+            raise typer.BadParameter(f"Symbolic links are not supported for flow file path.")
+
+        if file_path.suffix.lower() == ".py":
+            
+            # borrow code from python tool import to be able to load the script that holds the flow model
+
+            resolved_package_root = get_package_root(str(file_path.parent))
+            if resolved_package_root:
+                resolved_package_root = str(Path(resolved_package_root).absolute())
+                package_path = str(Path(resolved_package_root).parent.absolute())
+                package_folder = str(Path(resolved_package_root).stem)
+                sys.path.append(package_path)           # allows you to resolve non relative imports relative to the root of the module
+                sys.path.append(resolved_package_root)  # allows you to resolve relative imports in combination with import_module(..., package=...)
+                package = file_path_str.replace(resolved_package_root, '').replace('.py', '').replace('/', '.').replace('\\', '.')
+                if not path.isdir(resolved_package_root):
+                    raise typer.BadParameter(f"The provided package root is not a directory.")
+
+                elif not file_path_str.startswith(str(Path(resolved_package_root))):
+                    raise typer.BadParameter(f"The provided tool file path does not belong to the provided package root.")
+
+                temp_path = Path(file_path_str[len(str(Path(resolved_package_root))) + 1:])
+                if any([__supported_characters_pattern.match(x) is None for x in temp_path.parts[:-1]]):
+                    raise typer.BadParameter(f"Path to tool file contains unsupported characters. Only alphanumeric characters and underscores are allowed. Path: \"{temp_path}\"")
+            else:
+                package_folder = file_path.parent
+                package = file_path.stem
+                sys.path.append(str(package_folder))
+
+            module = importlib.import_module(package, package=package_folder)
+            
+            if resolved_package_root:
+                del sys.path[-1]
+            del sys.path[-1]
+
+            for _, obj in inspect.getmembers(module):
+                    
+                if not isinstance(obj, FlowWrapper):
+                    continue
+                
+                model = obj().to_json()
+                break
+
+        elif file_path.suffix.lower() == ".json":
+            with open(file) as f:
+                model = json.load(f)
+        else:
+            raise typer.BadParameter(f"Unknown file type.  Only python or json are supported.")
+
+
+    except typer.BadParameter as ex:
+        raise ex
+    
+    except Exception as e:
+        raise typer.BadParameter(f"Failed to load model from file {file}: {e}")
+    
+    return await import_flow_model(model)
+
+
 async def import_openapi_tool(file: str, connection_id: str) -> List[BaseTool]:
     tools = await create_openapi_json_tools_from_uri(file, connection_id)
     return tools
 
+def _get_kind_from_spec(spec: dict) -> ToolKind:
+    name = spec.get("name")
+    tool_binding = spec.get("binding")
+
+    if ToolKind.python in tool_binding:
+        return ToolKind.python
+    elif ToolKind.openapi in tool_binding:
+        return ToolKind.openapi
+    else:
+        logger.error(f"Could not determine 'kind' of tool '{name}'")
+        sys.exit(1) 
+
+def get_whl_in_registry(registry_url: str, version: str) -> str| None:
+    orchestrate_links = requests.get(registry_url).text
+    wheel_files = [x.group(1) for x in re.finditer( r'href="(.*\.whl).*"', orchestrate_links)]
+    wheel_file = next(filter(lambda x: f"{version}-py3-none-any.whl" in x, wheel_files), None)
+    return wheel_file
+
 class ToolsController:
     def __init__(self, tool_kind: ToolKind = None, file: str = None, requirements_file: str = None):
         self.client = None
@@ -335,7 +465,7 @@ class ToolsController:
         # Ensure app_id is a list
         if args.get("app_id") and isinstance(args.get("app_id"), str):
             args["app_id"] = [args.get("app_id")]
-    
+
         validate_params(kind=kind, **args)
 
         match kind:
@@ -356,6 +486,8 @@ class ToolsController:
                     connection = connections_client.get_draft_by_app_id(app_id=app_id)
                     connection_id = connection.connection_id
                 tools = asyncio.run(import_openapi_tool(file=args["file"], connection_id=connection_id))
+            case "flow":
+                tools = asyncio.run(import_flow_tool(file=args["file"]))
             case "skill":
                 tools = []
                 logger.warning("Skill Import not implemented yet")
@@ -390,7 +522,7 @@ class ToolsController:
 
             for tool in tools:
                 tool_binding = tool.__tool_spec__.binding
-                
+
                 connection_ids = []
 
                 if tool_binding is not None:
@@ -400,9 +532,11 @@ class ToolsController:
                         for conn in tool_binding.python.connections:
                             connection_ids.append(tool_binding.python.connections[conn])
                     elif tool_binding.mcp is not None and hasattr(tool_binding.mcp, "connections"):
-                        for conn in tool_binding.mcp.connections:
-                            connection_ids.append(tool_binding.mcp.connections[conn])
-
+                        if tool_binding.mcp.connections is None:
+                            connection_ids.append(None)
+                        else:
+                            for conn in tool_binding.mcp.connections:
+                                connection_ids.append(tool_binding.mcp.connections[conn])
 
                 app_ids = []
                 for connection_id in connection_ids:
@@ -463,7 +597,7 @@ class ToolsController:
                 if len(existing_tools) > 1:
                     logger.error(f"Multiple existing tools found with name '{tool.__tool_spec__.name}'. Failed to update tool")
                     sys.exit(1)
-                
+
                 if len(existing_tools) > 0:
                     existing_tool = existing_tools[0]
                     exist = True
@@ -518,12 +652,14 @@ class ToolsController:
                         if registry_type == RegistryType.LOCAL:
                             requirements.append(f"/packages/ibm_watsonx_orchestrate-0.6.0-py3-none-any.whl\n")
                         elif registry_type == RegistryType.PYPI:
+                            wheel_file = get_whl_in_registry(registry_url='https://pypi.org/simple/ibm-watsonx-orchestrate', version=version)
+                            if not wheel_file:
+                                logger.error(f"Could not find ibm-watsonx-orchestrate@{version} on https://pypi.org/project/ibm-watsonx-orchestrate")
+                                exit(1)
                             requirements.append(f"ibm-watsonx-orchestrate=={version}\n")
                         elif registry_type == RegistryType.TESTPYPI:
                             override_version = cfg.get(PYTHON_REGISTRY_HEADER, PYTHON_REGISTRY_TEST_PACKAGE_VERSION_OVERRIDE_OPT) or version
-                            orchestrate_links = requests.get('https://test.pypi.org/simple/ibm-watsonx-orchestrate').text
-                            wheel_files = [x.group(1) for x in re.finditer( r'href="(.*\.whl).*"', orchestrate_links)]
-                            wheel_file = next(filter(lambda x: f"{override_version}-py3-none-any.whl" in x, wheel_files), None)
+                            wheel_file = get_whl_in_registry(registry_url='https://test.pypi.org/simple/ibm-watsonx-orchestrate', version=override_version)
                             if not wheel_file:
                                 logger.error(f"Could not find ibm-watsonx-orchestrate@{override_version} on https://test.pypi.org/project/ibm-watsonx-orchestrate")
                                 exit(1)
@@ -569,7 +705,7 @@ class ToolsController:
             self.get_client().upload_tools_artifact(tool_id=tool_id, file_path=tool_artifact)
 
         logger.info(f"Tool '{tool.__tool_spec__.name}' updated successfully")
-    
+
     def remove_tool(self, name: str):
         try:
             client = self.get_client()
@@ -587,3 +723,51 @@ class ToolsController:
         except requests.HTTPError as e:
             logger.error(e.response.text)
             exit(1)
+
+    def download_tool(self, name: str) -> bytes | None:
+        tool_client = self.get_client()
+        draft_tools = tool_client.get_draft_by_name(tool_name=name)
+        if len(draft_tools) > 1:
+            logger.error(f"Multiple existing tools found with name '{name}'. Failed to get tool")
+            sys.exit(1)
+        if len(draft_tools) == 0:
+            logger.error(f"No tool named '{name}' found")
+            sys.exit(1)
+
+        draft_tool = draft_tools[0]
+        draft_tool_kind = _get_kind_from_spec(draft_tool)
+        
+        # TODO: Add openapi tool support
+        if draft_tool_kind != ToolKind.python:
+            logger.warning(f"Skipping '{name}', {draft_tool_kind.value} tools are currently unsupported by export")
+            return
+
+        tool_id = draft_tool.get("id")
+
+        tool_artifacts_bytes = tool_client.download_tools_artifact(tool_id=tool_id)
+        return tool_artifacts_bytes
+    
+    def export_tool(self, name: str, output_path: str) -> None:
+        
+        output_file = Path(output_path)
+        output_file_extension = output_file.suffix
+        if output_file_extension != ".zip":
+            logger.error(f"Output file must end with the extension '.zip'. Provided file '{output_path}' ends with '{output_file_extension}'")
+            sys.exit(1)
+        
+        logger.info(f"Exporting tool definition for '{name}' to '{output_path}'")
+
+        tool_artifact_bytes = self.download_tool(name)
+
+        if not tool_artifact_bytes:
+            return
+        
+        with zipfile.ZipFile(io.BytesIO(tool_artifact_bytes), "r") as zip_file_in, \
+            zipfile.ZipFile(output_path, 'w') as zip_file_out:
+            
+            for item in zip_file_in.infolist():
+                buffer = zip_file_in.read(item.filename)
+                if (item.filename != 'bundle-format'):
+                    zip_file_out.writestr(item, buffer)
+        
+        logger.info(f"Successfully exported tool definition for '{name}' to '{output_path}'")