deepset-mcp 0.0.4rc1__py3-none-any.whl → 0.0.5__py3-none-any.whl

deepset_mcp/tool_factory.py

@@ -6,459 +6,327 @@
 
 import functools
 import inspect
-import logging
-import os
+import re
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
-from enum import StrEnum
 from typing import Any
 
-from mcp.server.fastmcp import FastMCP
+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,
+from deepset_mcp.config import DEFAULT_CLIENT_HEADER, DOCS_SEARCH_TOOL_NAME
+from deepset_mcp.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig, WorkspaceMode
+from deepset_mcp.tool_registry import TOOL_REGISTRY
+from deepset_mcp.tools.tokonomics import (
+    ObjectStore,
+    RichExplorer,
+    explorable,
+    explorable_and_referenceable,
+    referenceable,
 )
 
-# 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.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 apply_custom_args(base_func: Callable[..., Any], config: ToolConfig) -> Callable[..., Any]:
+    """
+    Applies custom keyword arguments defined in the ToolConfig to a function.
 
-def are_docs_available() -> bool:
-    """Checks if documentation search is available."""
-    return bool(
-        os.environ.get("DEEPSET_DOCS_WORKSPACE", False)
-        and os.environ.get("DEEPSET_DOCS_PIPELINE_NAME", False)
-        and os.environ.get("DEEPSET_DOCS_API_KEY", False)
-    )
+    Removes the partially applied keyword arguments from the function's signature and docstring.
 
+    :param base_func: The function to apply custom keyword arguments to.
+    :param config: The ToolConfig for the function.
+    :returns: Function with custom arguments applied and updated signature/docstring.
+    """
+    if not config.custom_args:
+        return base_func
 
-EXPLORER = RichExplorer(store=STORE)
+    @functools.wraps(base_func)
+    async def func_with_custom_args(*args: Any, **kwargs: Any) -> Any:
+        # Create a partial function with the custom arguments bound.
+        partial_func = functools.partial(base_func, **(config.custom_args or {}))
+        # Await the result of the partial function call.
+        return await partial_func(**kwargs)
 
+    # Remove custom args from signature
+    original_sig = inspect.signature(base_func)
+    new_params = [p for name, p in original_sig.parameters.items() if name not in config.custom_args]
+    func_with_custom_args.__signature__ = original_sig.replace(parameters=new_params)  # type: ignore
 
-def get_from_object_store(object_id: str, path: str = "") -> str:
-    """Use this tool to fetch an object from the object store.
+    # Remove custom args from docstring.
+    func_with_custom_args.__doc__ = remove_params_from_docstring(base_func.__doc__, set(config.custom_args.keys()))
 
-    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).
+    return func_with_custom_args
 
-    :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.
+
+def remove_params_from_docstring(docstring: str | None, params_to_remove: set[str]) -> str:
+    """Removes specified parameters from a function's docstring.
+
+    :param docstring: The docstring to remove the parameters from.
+    :param params_to_remove: The set of parameters to remove.
+    :returns: The changed docstring.
     """
-    return EXPLORER.slice(obj_id=object_id, start=start, end=end, path=path)
+    if docstring is None:
+        return ""
+
+    for param_name in params_to_remove:
+        docstring = re.sub(
+            rf"^\s*:param\s+{re.escape(param_name)}.*?(?=^\s*:|^\s*$|\Z)",
+            "",
+            docstring,
+            flags=re.MULTILINE | re.DOTALL,
+        )
+
+    return "\n".join([line.rstrip() for line in docstring.strip().split("\n")])
 
 
-async def search_docs(query: str) -> str:
-    """Search the deepset platform documentation.
+def apply_workspace(
+    base_func: Callable[..., Any], config: ToolConfig, workspace_mode: WorkspaceMode, workspace: str | None = None
+) -> Callable[..., Any]:
+    """
+    Applies a deepset workspace to the function depending on the workspace mode and the ToolConfig.
 
-    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.
+    Removes the workspace argument from the function's signature and docstring if applied.
 
-    :param query: The search query to execute against the documentation.
-    :returns: The formatted search results from the documentation.
+    :param base_func: The function to apply workspace to.
+    :param config: The ToolConfig for the function.
+    :param workspace_mode: The WorkspaceMode for the function.
+    :param workspace: The workspace to use for static mode.
+    :returns: Function with workspace handling applied and updated signature/docstring.
+    :raises ValueError: If workspace is required but not available.
     """
-    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] | None = None  # For special cases like search_component_definition
-
-
-def get_workspace_from_env() -> str:
-    """Gets the workspace configured from environment variable."""
-    workspace = os.environ.get("DEEPSET_WORKSPACE")
-    if not workspace:
-        raise ValueError("DEEPSET_WORKSPACE environment variable not set")
-    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 create_enhanced_tool(
-    base_func: Callable[..., Any], config: ToolConfig, workspace_mode: WorkspaceMode, workspace: str | None = None
-) -> Callable[..., Awaitable[Any]]:
-    """Universal tool creator that handles client injection, workspace, and decorators.
+    if not config.needs_workspace:
+        return base_func
 
-    This function takes a base tool function and enhances it based on a configuration.
-    It can inject a `client`, manage a `workspace` parameter (either explicitly required
-    or implicitly provided from the environment), and apply memory-related decorators.
+    if workspace_mode == WorkspaceMode.STATIC:
 
-    It also supports partial application of custom arguments specified in the ToolConfig.
-    These arguments are bound to the function, and both the function signature and the
-    docstring are updated to hide these implementation details from the end user of the tool.
+        @functools.wraps(base_func)
+        async def workspace_wrapper(*args: Any, **kwargs: Any) -> Any:
+            return await base_func(*args, workspace=workspace, **kwargs)
 
-    All parameters in the final tool signature are converted to be keyword-only to enforce
-    explicit naming of arguments in tool calls.
+        # Remove workspace from signature
+        original_sig = inspect.signature(base_func)
+        new_params = [p for name, p in original_sig.parameters.items() if name != "workspace"]
+        workspace_wrapper.__signature__ = original_sig.replace(parameters=new_params)  # type: ignore
 
-    Args:
-        base_func: The base tool function.
-        config: Tool configuration specifying dependencies and custom arguments.
-        workspace_mode: How the workspace should be handled.
-        workspace: The workspace to use when using a static workspace.
+        # Remove workspace from docstring
+        workspace_wrapper.__doc__ = remove_params_from_docstring(base_func.__doc__, {"workspace"})
+
+        return workspace_wrapper
+    else:
+        # For dynamic mode, workspace is passed as parameter
+        return base_func
+
+
+def apply_memory(
+    base_func: Callable[..., Any], config: ToolConfig, store: ObjectStore | None = None
+) -> Callable[..., Any]:
+    """
+    Applies memory decorators to a function if requested in the ToolConfig.
 
-    Returns:
-        An enhanced, awaitable tool function with an updated signature and docstring.
+    :param base_func: The function to apply memory decorator to.
+    :param config: The ToolConfig for the function.
+    :param store: The ObjectStore instance to use
+    :returns: Function with memory decorators applied.
+    :raises ValueError: If an invalid memory type is specified.
     """
-    original_func = base_func
+    if config.memory_type == MemoryType.NO_MEMORY:
+        return base_func
 
-    # If custom arguments are provided, create a wrapper that applies them.
-    # This wrapper preserves the original function's metadata so that decorators work correctly.
-    func_to_decorate: Any
-    if config.custom_args:
+    if store is None:
+        raise ValueError("ObjectStore instance is required for memory decorators")
 
-        @functools.wraps(original_func)
-        async def func_with_custom_args(*args: Any, **kwargs: Any) -> Any:
-            # Create a partial function with the custom arguments bound.
-            partial_func = functools.partial(original_func, **(config.custom_args or {}))
-            # Await the result of the partial function call.
-            return await partial_func(**kwargs)
+    explorer = RichExplorer(store)
 
-        func_to_decorate = func_with_custom_args
+    if config.memory_type == MemoryType.EXPLORABLE:
+        return explorable(object_store=store, explorer=explorer)(base_func)
+    elif config.memory_type == MemoryType.REFERENCEABLE:
+        return referenceable(object_store=store, explorer=explorer)(base_func)
+    elif config.memory_type == MemoryType.BOTH:
+        return explorable_and_referenceable(object_store=store, explorer=explorer)(base_func)
     else:
-        func_to_decorate = original_func
-
-    # Apply memory-related decorators to the (potentially wrapped) function
-    decorated_func = func_to_decorate
-    if config.memory_type != MemoryType.NO_MEMORY:
-        store = STORE
-        explorer = RichExplorer(store)
-
-        if config.memory_type == MemoryType.EXPLORABLE:
-            decorated_func = explorable(object_store=store, explorer=explorer)(decorated_func)
-        elif config.memory_type == MemoryType.REFERENCEABLE:
-            decorated_func = referenceable(object_store=store, explorer=explorer)(decorated_func)
-        elif config.memory_type == MemoryType.BOTH:
-            decorated_func = explorable_and_referenceable(object_store=store, explorer=explorer)(decorated_func)
-
-    # Determine the parameters to remove from the original function's signature
-    params_to_remove: set[str] = set()
-    if config.custom_args:
-        params_to_remove.update(config.custom_args.keys())
-    if config.needs_client:
-        params_to_remove.add("client")
-    if config.needs_workspace and workspace_mode == WorkspaceMode.STATIC:
-        params_to_remove.add("workspace")
-
-    # Create the new signature from the original function
-    original_sig = inspect.signature(original_func)
-    final_params = [p for name, p in original_sig.parameters.items() if name not in params_to_remove]
-
-    # Convert all positional-or-keyword parameters to be keyword-only
-    keyword_only_params = [
-        p.replace(kind=inspect.Parameter.KEYWORD_ONLY) if p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD else p
-        for p in final_params
-    ]
-    new_sig = original_sig.replace(parameters=keyword_only_params)
-
-    # Create the final wrapper function that handles client/workspace injection
-    if config.needs_client:
-        if config.needs_workspace:
-            if workspace_mode == WorkspaceMode.STATIC:
-
-                async def workspace_environment_wrapper(**kwargs: Any) -> Any:
-                    ws = workspace or get_workspace_from_env()
-                    async with AsyncDeepsetClient(transport_config=DEFAULT_CLIENT_HEADER) as client:
-                        return await decorated_func(client=client, workspace=ws, **kwargs)
-
-                wrapper = workspace_environment_wrapper
-            else:  # DYNAMIC mode
-
-                async def workspace_explicit_wrapper(**kwargs: Any) -> Any:
-                    async with AsyncDeepsetClient(transport_config=DEFAULT_CLIENT_HEADER) as client:
-                        # The first argument is the workspace, which must be passed by keyword.
-                        return await decorated_func(client=client, **kwargs)
-
-                wrapper = workspace_explicit_wrapper
-        else:  # Client-only tools
-
-            async def client_only_wrapper(**kwargs: Any) -> Any:
-                async with AsyncDeepsetClient(transport_config=DEFAULT_CLIENT_HEADER) as client:
-                    return await decorated_func(client=client, **kwargs)
-
-            wrapper = client_only_wrapper
-    else:  # No injection needed
-        if inspect.iscoroutinefunction(decorated_func):
-
-            async def no_injection_wrapper(**kwargs: Any) -> Any:
-                return await decorated_func(**kwargs)
-
-            wrapper = no_injection_wrapper
-        else:
+        raise ValueError(f"Invalid memory type: {config.memory_type}")
 
-            @functools.wraps(decorated_func)
-            async def async_wrapper(**kwargs: Any) -> Any:
-                return decorated_func(**kwargs)
-
-            wrapper = async_wrapper
-
-    # Set metadata on the final wrapper
-    wrapper.__signature__ = new_sig  # type: ignore
-    wrapper.__name__ = original_func.__name__
-
-    # Process the docstring to remove injected and partially applied parameters
-    if original_func.__doc__:
-        import re
-
-        doc = original_func.__doc__
-        params_to_remove_from_doc = set()
-        if config.needs_client:
-            params_to_remove_from_doc.add("client")
-        if config.needs_workspace and workspace_mode == WorkspaceMode.STATIC:
-            params_to_remove_from_doc.add("workspace")
-        if config.custom_args:
-            params_to_remove_from_doc.update(config.custom_args.keys())
-
-        for param_name in params_to_remove_from_doc:
-            doc = re.sub(
-                rf"^\s*:param\s+{re.escape(param_name)}.*?(?=^\s*:|^\s*$|\Z)",
-                "",
-                doc,
-                flags=re.MULTILINE | re.DOTALL,
-            )
 
-        wrapper.__doc__ = "\n".join([line.rstrip() for line in doc.strip().split("\n")])
+def apply_client(
+    base_func: Callable[..., Any],
+    config: ToolConfig,
+    use_request_context: bool = True,
+    base_url: str | None = None,
+    api_key: str | None = None,
+) -> Callable[..., Any]:
+    """
+    Applies the deepset API client to a function.
+
+    Optionally collects the API key from the request context, when use_request_context is True.
+    Modifies the function's signature and docstring to remove the client argument.
+    Adds a 'ctx' argument to the signature if the request context is used.
+
+    :param base_func: The function to apply the client to.
+    :param config: The ToolConfig for the function.
+    :param use_request_context: Whether to collect the API key from the request context.
+    :param base_url: Base URL for the deepset API.
+    :param api_key: The API key to use.
+    :returns: Function with client injection applied and updated signature/docstring.
+    :raises ValueError: If API key cannot be extracted from request context.
+    """
+    if not config.needs_client:
+        return base_func
+
+    if use_request_context:
+
+        @functools.wraps(base_func)
+        async def client_wrapper_with_context(*args: Any, **kwargs: Any) -> Any:
+            ctx = kwargs.pop("ctx", None)
+            if not ctx:
+                raise ValueError("Context is required for client authentication")
+
+            api_key = ctx.request_context.request.headers.get("Authorization")
+            if not api_key:
+                raise ValueError("No Authorization header found in request context")
+
+            api_key = api_key.replace("Bearer ", "")
+
+            if not api_key:
+                raise ValueError("API key cannot be empty")
+
+            client_kwargs: dict[str, Any] = {"transport_config": DEFAULT_CLIENT_HEADER, "api_key": api_key}
+            if base_url:
+                client_kwargs["base_url"] = base_url
+            async with AsyncDeepsetClient(**client_kwargs) as client:
+                return await base_func(*args, client=client, **kwargs)
+
+        # Remove client from signature and add ctx
+        original_sig = inspect.signature(base_func)
+        new_params = [p for name, p in original_sig.parameters.items() if name != "client"]
+        ctx_param = inspect.Parameter(name="ctx", kind=inspect.Parameter.KEYWORD_ONLY, annotation=Context)
+        new_params.append(ctx_param)
+        client_wrapper_with_context.__signature__ = original_sig.replace(parameters=new_params)  # type: ignore
+        client_wrapper_with_context.__doc__ = remove_params_from_docstring(base_func.__doc__, {"client"})
+
+        return client_wrapper_with_context
     else:
-        wrapper.__doc__ = original_func.__doc__
 
-    return wrapper
+        @functools.wraps(base_func)
+        async def client_wrapper_without_context(*args: Any, **kwargs: Any) -> Any:
+            client_kwargs: dict[str, Any] = {"transport_config": DEFAULT_CLIENT_HEADER, "api_key": api_key}
+            if base_url:
+                client_kwargs["base_url"] = base_url
+            async with AsyncDeepsetClient(**client_kwargs) as client:
+                return await base_func(*args, client=client, **kwargs)
+
+        # Remove client from signature
+        original_sig = inspect.signature(base_func)
+        new_params = [p for name, p in original_sig.parameters.items() if name != "client"]
+        client_wrapper_without_context.__signature__ = original_sig.replace(parameters=new_params)  # type: ignore
+
+        # Remove client from docstring
+        client_wrapper_without_context.__doc__ = remove_params_from_docstring(base_func.__doc__, {"client"})
+
+        return client_wrapper_without_context
+
+
+def build_tool(
+    base_func: Callable[..., Any],
+    config: ToolConfig,
+    workspace_mode: WorkspaceMode,
+    api_key: str | None = None,
+    workspace: str | None = None,
+    use_request_context: bool = True,
+    base_url: str | None = None,
+    object_store: ObjectStore | None = None,
+) -> Callable[..., Awaitable[Any]]:
+    """
+    Universal tool creator that handles client injection, workspace, and decorators.
+
+    This function takes a base tool function and enhances it based on the tool's configuration.
+
+    :param base_func: The base tool function.
+    :param config: Tool configuration specifying dependencies and custom arguments.
+    :param workspace_mode: How the workspace should be handled.
+    :param api_key: The deepset API key to use.
+    :param workspace: The workspace to use when using a static workspace.
+    :param use_request_context: Whether to collect the API key from the request context.
+    :param base_url: Base URL for the deepset API.
+    :param object_store: The ObjectStore instance to use for memory decorators.
+    :returns: An enhanced, awaitable tool function with an updated signature and docstring.
+    """
+    enhanced_func = base_func
+
+    # Apply custom arguments first
+    enhanced_func = apply_custom_args(enhanced_func, config)
+
+    # Apply memory decorators with the provided store
+    enhanced_func = apply_memory(enhanced_func, config, object_store)
+
+    # Apply workspace handling
+    enhanced_func = apply_workspace(enhanced_func, config, workspace_mode, workspace)
+
+    # Apply client injection (adds ctx parameter if needed)
+    enhanced_func = apply_client(
+        enhanced_func, config, use_request_context=use_request_context, base_url=base_url, api_key=api_key
+    )
+
+    # Create final async wrapper if needed
+    if not inspect.iscoroutinefunction(enhanced_func):
+
+        @functools.wraps(enhanced_func)
+        async def async_wrapper(**kwargs: Any) -> Any:
+            return enhanced_func(**kwargs)
+
+        # Copy over the signature from the enhanced function
+        async_wrapper.__signature__ = inspect.signature(enhanced_func)  # type: ignore
+        return async_wrapper
+
+    enhanced_func.__name__ = base_func.__name__
+
+    return enhanced_func
 
 
 def register_tools(
-    mcp: FastMCP, workspace_mode: WorkspaceMode, workspace: str | None = None, tool_names: set[str] | None = None
+    mcp_server_instance: FastMCP,
+    workspace_mode: WorkspaceMode,
+    api_key: str | None = None,
+    workspace: str | None = None,
+    tool_names: set[str] | None = None,
+    get_api_key_from_authorization_header: bool = True,
+    docs_config: DeepsetDocsConfig | None = None,
+    base_url: str | None = None,
+    object_store: ObjectStore | None = None,
 ) -> None:
     """Register tools with unified configuration.
 
     Args:
-        mcp: FastMCP server instance
+        mcp_server_instance: FastMCP server instance
         workspace_mode: How workspace should be handled
-        workspace: Workspace to use for environment mode (if None, reads from env)
+        api_key: An api key for the deepset AI platform; only needs to be provided when not read from request context.
+        workspace: Workspace to use; only needs to be provided if using a static workspace.
         tool_names: Set of tool names to register (if None, registers all tools)
+        get_api_key_from_authorization_header: Whether to use request context to retrieve an API key for tool execution.
+        docs_config: Configuration for the deepset documentation search tool.
+        base_url: Base URL for the deepset API.
+        object_store: The ObjectStore instance to use for memory decorators.
     """
-    # Check if docs search is available
-    docs_available = are_docs_available()
+    if api_key is None and not get_api_key_from_authorization_header:
+        raise ValueError(
+            "'api_key' cannot be 'None' when 'use_request_context' is False. "
+            "Either pass 'api_key' or 'use_request_context'."
+        )
+
+    if workspace_mode == WorkspaceMode.STATIC and workspace is None:
+        raise ValueError(
+            "'workspace_mode' set to 'static' but no workspace provided. "
+            "You need to set a deepset workspace name as 'workspace'."
+        )
+
+    if docs_config is None and tool_names is None:
+        raise ValueError(
+            f"'docs_config' cannot be None when requesting to register all tools. "
+            f"Either pass 'docs_config' or disable the '{DOCS_SEARCH_TOOL_NAME}' tool."
+        )
+
+    if docs_config is None and tool_names is not None and DOCS_SEARCH_TOOL_NAME in tool_names:
+        raise ValueError(
+            f"Requested to register '{DOCS_SEARCH_TOOL_NAME}' tool but 'docs_config' is 'None'. "
+            f"Provide a valid 'docs_config' to register this tool."
+        )
 
     # Validate tool names if provided
     if tool_names is not None:
@@ -469,30 +337,35 @@ def register_tools(
             sorted_all = sorted(all_tools)
             raise ValueError(f"Unknown tools: {', '.join(sorted_invalid)}\nAvailable tools: {', '.join(sorted_all)}")
 
-        # Warn if search_docs was requested but config is missing
-        if "search_docs" in tool_names and not docs_available:
-            logging.warning(
-                "Documentation search tool requested but not available. To enable, set the DEEPSET_DOCS_SHARE_URL "
-                "environment variable."
-            )
-
         tools_to_register = tool_names.copy()
     else:
         tools_to_register = set(TOOL_REGISTRY.keys())
 
-        # Warn if search_docs would be skipped in "all tools" mode
-        if not docs_available:
-            logging.warning(
-                "Documentation search tool not enabled. To enable, set the DEEPSET_DOCS_SHARE_URL environment variable."
-            )
-
-    # Remove search_docs if config is not available
-    if not docs_available:
-        tools_to_register.discard("search_docs")
-
     for tool_name in tools_to_register:
         base_func, config = TOOL_REGISTRY[tool_name]
-        # Create enhanced tool
-        enhanced_tool = create_enhanced_tool(base_func, config, workspace_mode, workspace)
 
-        mcp.add_tool(enhanced_tool, name=tool_name, structured_output=False)
+        if tool_name == DOCS_SEARCH_TOOL_NAME:
+            # search_docs is a special tool.
+            # base_func is a factory function.
+            # We configure with the docs_config to get the actual tool function.
+            enhanced_tool = base_func(config=docs_config)
+        elif tool_name in ("get_from_object_store", "get_slice_from_object_store"):
+            # ObjectStore tools are factory functions that need an explorer created from the store
+            if object_store is None:
+                raise ValueError(f"ObjectStore instance is required for {tool_name}")
+
+            explorer = RichExplorer(store=object_store)
+            enhanced_tool = base_func(explorer=explorer)
+        else:
+            enhanced_tool = build_tool(
+                base_func=base_func,
+                config=config,
+                workspace_mode=workspace_mode,
+                workspace=workspace,
+                use_request_context=get_api_key_from_authorization_header,
+                base_url=base_url,
+                object_store=object_store,
+                api_key=api_key,
+            )
+
+        mcp_server_instance.add_tool(enhanced_tool, name=tool_name)