refactor: remove workspace mode (#193)

mathislucka · web-flow · commit d42ccad77d5e · 2025-08-25T11:50:34.000+02:00
diff --git a/src/deepset_mcp/__init__.py b/src/deepset_mcp/__init__.py
@@ -3,8 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from deepset_mcp.config import DEEPSET_DOCS_DEFAULT_SHARE_URL
-from deepset_mcp.server import configure_mcp_server
-from deepset_mcp.tool_models import WorkspaceMode
-from deepset_mcp.tool_registry import ALL_DEEPSET_TOOLS
+from deepset_mcp.mcp.server import configure_mcp_server
+from deepset_mcp.mcp.tool_models import WorkspaceMode
+from deepset_mcp.mcp.tool_registry import ALL_DEEPSET_TOOLS
 
 __all__ = ["configure_mcp_server", "WorkspaceMode", "ALL_DEEPSET_TOOLS", "DEEPSET_DOCS_DEFAULT_SHARE_URL"]
diff --git a/src/deepset_mcp/main.py b/src/deepset_mcp/main.py
@@ -11,9 +11,9 @@
 from mcp.server.fastmcp import FastMCP
 
 from deepset_mcp.config import DEEPSET_DOCS_DEFAULT_SHARE_URL, DOCS_SEARCH_TOOL_NAME
-from deepset_mcp.server import configure_mcp_server
-from deepset_mcp.tool_models import WorkspaceMode
-from deepset_mcp.tool_registry import TOOL_REGISTRY
+from deepset_mcp.mcp.server import configure_mcp_server
+from deepset_mcp.mcp.tool_models import WorkspaceMode
+from deepset_mcp.mcp.tool_registry import TOOL_REGISTRY
 
 
 class TransportEnum(StrEnum):
diff --git a/src/deepset_mcp/mcp/__init__.py b/src/deepset_mcp/mcp/__init__.py
@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .server import configure_mcp_server
+from .store import initialize_or_get_initialized_store
+from .tool_factory import build_tool
+from .tool_models import ToolConfig, WorkspaceMode
+
+__all__ = ["configure_mcp_server", "build_tool", "ToolConfig", "WorkspaceMode", "initialize_or_get_initialized_store"]
diff --git a/src/deepset_mcp/mcp/server.py b/src/deepset_mcp/mcp/server.py
@@ -10,10 +10,10 @@
 
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.config import DEEPSET_DOCS_DEFAULT_SHARE_URL
-from deepset_mcp.store import initialize_store
-from deepset_mcp.tool_factory import register_tools
-from deepset_mcp.tool_models import DeepsetDocsConfig, WorkspaceMode
-from deepset_mcp.tool_registry import TOOL_REGISTRY
+from deepset_mcp.mcp.store import initialize_or_get_initialized_store
+from deepset_mcp.mcp.tool_factory import register_tools
+from deepset_mcp.mcp.tool_models import DeepsetDocsConfig, WorkspaceMode
+from deepset_mcp.mcp.tool_registry import TOOL_REGISTRY
 
 
 def configure_mcp_server(
@@ -73,7 +73,9 @@ def configure_mcp_server(
     docs_config = DeepsetDocsConfig(api_key=api_key_docs, workspace_name=workspace_name, pipeline_name=pipeline_name)
 
     # Initialize the store before registering tools
-    store = initialize_store(backend=object_store_backend, redis_url=object_store_redis_url, ttl=object_store_ttl)
+    store = initialize_or_get_initialized_store(
+        backend=object_store_backend, redis_url=object_store_redis_url, ttl=object_store_ttl
+    )
 
     register_tools(
         mcp_server_instance=mcp_server_instance,
diff --git a/src/deepset_mcp/mcp/store.py b/src/deepset_mcp/mcp/store.py
@@ -30,12 +30,12 @@ def create_redis_backend(url: str) -> ObjectStoreBackend:
 
 
 @functools.lru_cache(maxsize=1)
-def initialize_store(
+def initialize_or_get_initialized_store(
     backend: str = "memory",
     redis_url: str | None = None,
     ttl: int = 600,
 ) -> ObjectStore:
-    """Initialize the object store.
+    """Initializes the object store or gets an existing object store instance if it was initialized before.
 
     :param backend: Backend type ('memory' or 'redis')
     :param redis_url: Redis connection URL (required if backend='redis')
diff --git a/src/deepset_mcp/mcp/tool_factory.py b/src/deepset_mcp/mcp/tool_factory.py
@@ -14,15 +14,15 @@
 
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.config import DEFAULT_CLIENT_HEADER, DOCS_SEARCH_TOOL_NAME
+from deepset_mcp.mcp.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig, WorkspaceMode
+from deepset_mcp.mcp.tool_registry import TOOL_REGISTRY
 from deepset_mcp.tokonomics import (
     ObjectStore,
     RichExplorer,
     explorable,
     explorable_and_referenceable,
     referenceable,
 )
-from deepset_mcp.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig, WorkspaceMode
-from deepset_mcp.tool_registry import TOOL_REGISTRY
 
 
 def apply_custom_args(base_func: Callable[..., Any], config: ToolConfig) -> Callable[..., Any]:
@@ -78,41 +78,35 @@ def remove_params_from_docstring(docstring: str | None, params_to_remove: set[st
 
 
 def apply_workspace(
-    base_func: Callable[..., Any], config: ToolConfig, workspace_mode: WorkspaceMode, workspace: str | None = None
+    base_func: Callable[..., Any], config: ToolConfig, workspace: str | None = None
 ) -> Callable[..., Any]:
     """
-    Applies a deepset workspace to the function depending on the workspace mode and the ToolConfig.
+    Applies a deepset workspace to the function depending on the ToolConfig.
 
     Removes the workspace argument from the function's signature and docstring if applied.
 
     :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.
+    :param workspace: The workspace to use.
     :returns: Function with workspace handling applied and updated signature/docstring.
     :raises ValueError: If workspace is required but not available.
     """
-    if not config.needs_workspace:
+    if not config.needs_workspace or not workspace:
         return base_func
 
-    if workspace_mode == WorkspaceMode.STATIC:
-
-        @functools.wraps(base_func)
-        async def workspace_wrapper(*args: Any, **kwargs: Any) -> Any:
-            return await base_func(*args, workspace=workspace, **kwargs)
+    @functools.wraps(base_func)
+    async def workspace_wrapper(*args: Any, **kwargs: Any) -> Any:
+        return await base_func(*args, workspace=workspace, **kwargs)
 
-        # 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
+    # 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
 
-        # Remove workspace from docstring
-        workspace_wrapper.__doc__ = remove_params_from_docstring(base_func.__doc__, {"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
+    return workspace_wrapper
 
 
 def apply_memory(
@@ -226,7 +220,6 @@ async def client_wrapper_without_context(*args: Any, **kwargs: Any) -> Any:
 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,
@@ -240,7 +233,6 @@ def build_tool(
 
     :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.
@@ -257,7 +249,7 @@ def build_tool(
     enhanced_func = apply_memory(enhanced_func, config, object_store)
 
     # Apply workspace handling
-    enhanced_func = apply_workspace(enhanced_func, config, workspace_mode, workspace)
+    enhanced_func = apply_workspace(base_func=enhanced_func, config=config, workspace=workspace)
 
     # Apply client injection (adds ctx parameter if needed)
     enhanced_func = apply_client(
@@ -360,7 +352,6 @@ def register_tools(
             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,
diff --git a/src/deepset_mcp/mcp/tool_models.py b/src/deepset_mcp/mcp/tool_models.py
diff --git a/src/deepset_mcp/mcp/tool_registry.py b/src/deepset_mcp/mcp/tool_registry.py
@@ -8,7 +8,7 @@
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.config import DEFAULT_CLIENT_HEADER, DOCS_SEARCH_TOOL_NAME
 from deepset_mcp.initialize_embedding_model import get_initialized_model
-from deepset_mcp.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig
+from deepset_mcp.mcp.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig
 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,
diff --git a/test/unit/test_server_base_url.py b/test/unit/test_server_base_url.py
@@ -6,14 +6,14 @@
 
 from unittest.mock import MagicMock, patch
 
-from deepset_mcp.server import configure_mcp_server
-from deepset_mcp.tool_models import WorkspaceMode
+from deepset_mcp.mcp.server import configure_mcp_server
+from deepset_mcp.mcp.tool_models import WorkspaceMode
 
 
 class TestConfigureMcpServerBaseUrl:
     """Test the configure_mcp_server function with base_url parameter."""
 
-    @patch("deepset_mcp.server.register_tools")
+    @patch("deepset_mcp.mcp.server.register_tools")
     def test_configure_mcp_server_passes_base_url(self, mock_register_tools: MagicMock) -> None:
         """Test that configure_mcp_server passes base_url to register_tools."""
         mock_server = MagicMock()
@@ -33,7 +33,7 @@ def test_configure_mcp_server_passes_base_url(self, mock_register_tools: MagicMo
         call_args = mock_register_tools.call_args
         assert call_args[1]["base_url"] == custom_url
 
-    @patch("deepset_mcp.server.register_tools")
+    @patch("deepset_mcp.mcp.server.register_tools")
     def test_configure_mcp_server_without_base_url(self, mock_register_tools: MagicMock) -> None:
         """Test that configure_mcp_server works without base_url."""
         mock_server = MagicMock()
diff --git a/test/unit/test_store.py b/test/unit/test_store.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from deepset_mcp.store import create_redis_backend, initialize_store
+from deepset_mcp.mcp.store import create_redis_backend, initialize_or_get_initialized_store
 from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore
 
 
@@ -15,15 +15,15 @@ class TestStoreInitialization:
 
     def test_initialize_store_memory_backend(self) -> None:
         """Test initializing store with memory backend."""
-        store = initialize_store(backend="memory", ttl=1800)
+        store = initialize_or_get_initialized_store(backend="memory", ttl=1800)
 
         assert isinstance(store, ObjectStore)
         assert isinstance(store._backend, InMemoryBackend)
         assert store._ttl == 1800
 
     def test_initialize_store_default_backend(self) -> None:
         """Test initializing store with default backend."""
-        store = initialize_store()
+        store = initialize_or_get_initialized_store()
 
         assert isinstance(store, ObjectStore)
         assert isinstance(store._backend, InMemoryBackend)
@@ -35,7 +35,7 @@ def test_initialize_store_redis_backend_success(self) -> None:
         mock_redis_client.ping.return_value = True
 
         with patch("redis.from_url", return_value=mock_redis_client):
-            store = initialize_store(backend="redis", redis_url="redis://localhost:6379", ttl=7200)
+            store = initialize_or_get_initialized_store(backend="redis", redis_url="redis://localhost:6379", ttl=7200)
 
             assert isinstance(store, ObjectStore)
             assert store._ttl == 7200
@@ -44,7 +44,7 @@ def test_initialize_store_redis_backend_success(self) -> None:
     def test_initialize_store_redis_backend_no_url(self) -> None:
         """Test initializing store with Redis backend but no URL."""
         with pytest.raises(ValueError, match="redis_url.*is None"):
-            initialize_store(backend="redis", redis_url=None)
+            initialize_or_get_initialized_store(backend="redis", redis_url=None)
 
     def test_initialize_store_redis_backend_connection_failure(self) -> None:
         """Test initializing store with Redis backend connection failure."""
@@ -53,7 +53,7 @@ def test_initialize_store_redis_backend_connection_failure(self) -> None:
 
         with patch("redis.from_url", return_value=mock_redis_client):
             with pytest.raises(Exception, match="Connection failed"):
-                initialize_store(backend="redis", redis_url="redis://localhost:6379")
+                initialize_or_get_initialized_store(backend="redis", redis_url="redis://localhost:6379")
 
     def test_create_redis_backend_success(self) -> None:
         """Test creating Redis backend successfully."""
@@ -85,21 +85,21 @@ def test_create_redis_backend_connection_failure(self) -> None:
     def test_initialize_store_caching(self) -> None:
         """Test that initialize_store uses caching."""
         # Clear any existing cache
-        initialize_store.cache_clear()
+        initialize_or_get_initialized_store.cache_clear()
 
-        store1 = initialize_store(backend="memory", ttl=1800)
-        store2 = initialize_store(backend="memory", ttl=1800)
+        store1 = initialize_or_get_initialized_store(backend="memory", ttl=1800)
+        store2 = initialize_or_get_initialized_store(backend="memory", ttl=1800)
 
         # Should return the same instance due to caching
         assert store1 is store2
 
     def test_initialize_store_different_params_different_instances(self) -> None:
         """Test that different parameters create different instances."""
         # Clear any existing cache
-        initialize_store.cache_clear()
+        initialize_or_get_initialized_store.cache_clear()
 
-        store1 = initialize_store(backend="memory", ttl=1800)
-        store2 = initialize_store(backend="memory", ttl=3600)
+        store1 = initialize_or_get_initialized_store(backend="memory", ttl=1800)
+        store2 = initialize_or_get_initialized_store(backend="memory", ttl=3600)
 
         # Should return different instances due to different parameters
         assert store1 is not store2
@@ -108,7 +108,7 @@ def test_initialize_store_different_params_different_instances(self) -> None:
 
     def test_initialize_store_unknown_backend(self) -> None:
         """Test initializing store with unknown backend defaults to memory."""
-        store = initialize_store(backend="unknown", ttl=1800)
+        store = initialize_or_get_initialized_store(backend="unknown", ttl=1800)
 
         assert isinstance(store, ObjectStore)
         assert isinstance(store._backend, InMemoryBackend)
diff --git a/test/unit/test_tool_factory.py b/test/unit/test_tool_factory.py
