refactor: move tool registry (#160)

* refactor: move tool registry

* chore: license headers

Author: mathislucka

src/deepset_mcp/main.py

@@ -14,7 +14,8 @@
 
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.config import DEEPSET_DOCS_DEFAULT_SHARE_URL
-from deepset_mcp.tool_factory import WorkspaceMode, register_tools
+from deepset_mcp.tool_factory import register_tools
+from deepset_mcp.tool_models import WorkspaceMode
 
 # Initialize MCP Server
 mcp = FastMCP("Deepset Cloud MCP", settings={"log_level": "ERROR"})
@@ -140,7 +141,7 @@ def main() -> None:
 
     # Handle --list-tools flag early
     if args.list_tools:
-        from deepset_mcp.tool_factory import TOOL_REGISTRY
+        from deepset_mcp.tool_registry import TOOL_REGISTRY
 
         print("Available tools:")
         for tool_name in sorted(TOOL_REGISTRY.keys()):

src/deepset_mcp/tool_factory.py

@@ -10,63 +10,16 @@
 import os
 import re
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from enum import StrEnum
 from typing import Any
 
 from mcp.server.fastmcp import Context, FastMCP
 
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.config import DEFAULT_CLIENT_HEADER
-from deepset_mcp.initialize_embedding_model import get_initialized_model
 from deepset_mcp.store import STORE
-from deepset_mcp.tools.custom_components import (
-    get_latest_custom_component_installation_logs as get_latest_custom_component_installation_logs_tool,
-    list_custom_component_installations as list_custom_component_installations_tool,
-)
-from deepset_mcp.tools.doc_search import (
-    search_docs as search_docs_tool,
-)
-from deepset_mcp.tools.haystack_service import (
-    get_component_definition as get_component_definition_tool,
-    get_custom_components as get_custom_components_tool,
-    list_component_families as list_component_families_tool,
-    search_component_definition as search_component_definition_tool,
-)
-from deepset_mcp.tools.indexes import (
-    create_index as create_index_tool,
-    deploy_index as deploy_index_tool,
-    get_index as get_index_tool,
-    list_indexes as list_indexes_tool,
-    update_index as update_index_tool,
-)
-
-# Import all tool functions
-from deepset_mcp.tools.pipeline import (
-    create_pipeline as create_pipeline_tool,
-    deploy_pipeline as deploy_pipeline_tool,
-    get_pipeline as get_pipeline_tool,
-    get_pipeline_logs as get_pipeline_logs_tool,
-    list_pipelines as list_pipelines_tool,
-    search_pipeline as search_pipeline_tool,
-    update_pipeline as update_pipeline_tool,
-    validate_pipeline as validate_pipeline_tool,
-)
-from deepset_mcp.tools.pipeline_template import (
-    get_template as get_pipeline_template_tool,
-    list_templates as list_pipeline_templates_tool,
-    search_templates as search_pipeline_templates_tool,
-)
-from deepset_mcp.tools.secrets import (
-    get_secret as get_secret_tool,
-    list_secrets as list_secrets_tool,
-)
+from deepset_mcp.tool_models import MemoryType, ToolConfig, WorkspaceMode
+from deepset_mcp.tool_registry import TOOL_REGISTRY
 from deepset_mcp.tools.tokonomics import RichExplorer, explorable, explorable_and_referenceable, referenceable
-from deepset_mcp.tools.workspace import (
-    create_workspace as create_workspace_tool,
-    get_workspace as get_workspace_tool,
-    list_workspaces as list_workspaces_tool,
-)
 
 
 def are_docs_available() -> bool:
@@ -78,89 +31,6 @@ def are_docs_available() -> bool:
     )
 
 
-EXPLORER = RichExplorer(store=STORE)
-
-
-def get_from_object_store(object_id: str, path: str = "") -> str:
-    """Use this tool to fetch an object from the object store.
-
-    You can fetch a specific object by using the object's id (e.g. `@obj_001`).
-    You can also fetch any nested path by using the path-parameter
-        (e.g. `{"object_id": "@obj_001", "path": "user_info.given_name"}`
-        -> returns the content at obj.user_info.given_name).
-
-    :param object_id: The id of the object to fetch in the format `@obj_001`.
-    :param path: The path of the object to fetch in the format of `access.to.attr` or `["access"]["to"]["attr"]`.
-    """
-    return EXPLORER.explore(obj_id=object_id, path=path)
-
-
-def get_slice_from_object_store(
-    object_id: str,
-    start: int = 0,
-    end: int | None = None,
-    path: str = "",
-) -> str:
-    """Extract a slice from a string or list object that is stored in the object store.
-
-    :param object_id: Identifier of the object.
-    :param start: Start index for slicing.
-    :param end: End index for slicing (optional - leave empty to get slice from start to end of sequence).
-    :param path: Navigation path to object to slice (optional).
-    :return: String representation of the slice.
-    """
-    return EXPLORER.slice(obj_id=object_id, start=start, end=end, path=path)
-
-
-async def search_docs(query: str) -> str:
-    """Search the deepset platform documentation.
-
-    This tool allows you to search through deepset's official documentation to find
-    information about features, API usage, best practices, and troubleshooting guides.
-    Use this when you need to look up specific deepset functionality or help users
-    understand how to use deepset features.
-
-    :param query: The search query to execute against the documentation.
-    :returns: The formatted search results from the documentation.
-    """
-    async with AsyncDeepsetClient(
-        api_key=os.environ["DEEPSET_DOCS_API_KEY"], transport_config=DEFAULT_CLIENT_HEADER
-    ) as client:
-        response = await search_docs_tool(
-            client=client,
-            workspace=os.environ["DEEPSET_DOCS_WORKSPACE"],
-            pipeline_name=os.environ["DEEPSET_DOCS_PIPELINE_NAME"],
-            query=query,
-        )
-    return response
-
-
-class WorkspaceMode(StrEnum):
-    """Configuration for how workspace is provided to tools."""
-
-    STATIC = "static"  # workspace from env, no parameter in tool signature
-    DYNAMIC = "dynamic"  # workspace as required parameter in tool signature
-
-
-class MemoryType(StrEnum):
-    """Configuration for how memory is provided to tools."""
-
-    EXPLORABLE = "explorable"
-    REFERENCEABLE = "referenceable"
-    BOTH = "both"
-    NO_MEMORY = "no_memory"
-
-
-@dataclass
-class ToolConfig:
-    """Configuration for tool registration."""
-
-    needs_client: bool = False
-    needs_workspace: bool = False
-    memory_type: MemoryType = MemoryType.NO_MEMORY
-    custom_args: dict[str, Any] = field(default_factory=dict)
-
-
 def get_workspace_from_env() -> str:
     """Gets the workspace configured from environment variable."""
     workspace = os.environ.get("DEEPSET_WORKSPACE")
@@ -169,135 +39,6 @@ def get_workspace_from_env() -> str:
     return workspace
 
 
-TOOL_REGISTRY: dict[str, tuple[Callable[..., Any], ToolConfig]] = {
-    # Workspace tools
-    "list_pipelines": (
-        list_pipelines_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "create_pipeline": (
-        create_pipeline_tool,
-        ToolConfig(
-            needs_client=True,
-            needs_workspace=True,
-            memory_type=MemoryType.BOTH,
-            custom_args={"skip_validation_errors": True},
-        ),
-    ),
-    "update_pipeline": (
-        update_pipeline_tool,
-        ToolConfig(
-            needs_client=True,
-            needs_workspace=True,
-            memory_type=MemoryType.BOTH,
-            custom_args={"skip_validation_errors": True},
-        ),
-    ),
-    "get_pipeline": (
-        get_pipeline_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "deploy_pipeline": (
-        deploy_pipeline_tool,
-        ToolConfig(
-            needs_client=True,
-            needs_workspace=True,
-            memory_type=MemoryType.EXPLORABLE,
-            custom_args={"wait_for_deployment": True, "timeout_seconds": 600, "poll_interval": 5},
-        ),
-    ),
-    "validate_pipeline": (
-        validate_pipeline_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
-    ),
-    "get_pipeline_logs": (
-        get_pipeline_logs_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "search_pipeline": (
-        search_pipeline_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "list_indexes": (
-        list_indexes_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "get_index": (
-        get_index_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "create_index": (
-        create_index_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
-    ),
-    "update_index": (
-        update_index_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
-    ),
-    "deploy_index": (
-        deploy_index_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "list_templates": (
-        list_pipeline_templates_tool,
-        ToolConfig(
-            needs_client=True,
-            needs_workspace=True,
-            memory_type=MemoryType.EXPLORABLE,
-            custom_args={"field": "created_at", "order": "DESC", "limit": 100},
-        ),
-    ),
-    "get_template": (
-        get_pipeline_template_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "search_templates": (
-        search_pipeline_templates_tool,
-        ToolConfig(
-            needs_client=True,
-            needs_workspace=True,
-            memory_type=MemoryType.EXPLORABLE,
-            custom_args={"model": get_initialized_model()},
-        ),
-    ),
-    "list_custom_component_installations": (
-        list_custom_component_installations_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "get_latest_custom_component_installation_logs": (
-        get_latest_custom_component_installation_logs_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    # Non-workspace tools
-    "list_component_families": (
-        list_component_families_tool,
-        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "get_component_definition": (
-        get_component_definition_tool,
-        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "search_component_definitions": (
-        search_component_definition_tool,
-        ToolConfig(
-            needs_client=True, memory_type=MemoryType.EXPLORABLE, custom_args={"model": get_initialized_model()}
-        ),
-    ),
-    "get_custom_components": (
-        get_custom_components_tool,
-        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE),
-    ),
-    "list_secrets": (list_secrets_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
-    "get_secret": (get_secret_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
-    "list_workspaces": (list_workspaces_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
-    "get_workspace": (get_workspace_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
-    "create_workspace": (create_workspace_tool, ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE)),
-    "get_from_object_store": (get_from_object_store, ToolConfig(memory_type=MemoryType.NO_MEMORY)),
-    "get_slice_from_object_store": (get_slice_from_object_store, ToolConfig(memory_type=MemoryType.NO_MEMORY)),
-    "search_docs": (search_docs, ToolConfig(memory_type=MemoryType.NO_MEMORY)),
-}
-
-
 def apply_custom_args(base_func: Callable[..., Any], config: ToolConfig) -> Callable[..., Any]:
     """
     Applies custom keyword arguments defined in the ToolConfig to a function.

src/deepset_mcp/tool_models.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any
+
+
+class MemoryType(StrEnum):
+    """Configuration for how memory is provided to tools."""
+
+    EXPLORABLE = "explorable"
+    REFERENCEABLE = "referenceable"
+    BOTH = "both"
+    NO_MEMORY = "no_memory"
+
+
+@dataclass
+class ToolConfig:
+    """Configuration for tool registration."""
+
+    needs_client: bool = False
+    needs_workspace: bool = False
+    memory_type: MemoryType = MemoryType.NO_MEMORY
+    custom_args: dict[str, Any] = field(default_factory=dict)
+
+
+class WorkspaceMode(StrEnum):
+    """Configuration for how workspace is provided to tools."""
+
+    STATIC = "static"  # workspace from env, no parameter in tool signature
+    DYNAMIC = "dynamic"  # workspace as required parameter in tool signature