diff --git a/src/deepset_mcp/__init__.py b/src/deepset_mcp/__init__.py index ca06e90..0c26002 100644 --- 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 index 7e0baa2..e6be415 100644 --- 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 new file mode 100644 index 0000000..3b92e4a --- /dev/null +++ b/src/deepset_mcp/mcp/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: 2025-present deepset GmbH +# +# 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/server.py b/src/deepset_mcp/mcp/server.py similarity index 94% rename from src/deepset_mcp/server.py rename to src/deepset_mcp/mcp/server.py index a0ee19f..6469c54 100644 --- a/src/deepset_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/store.py b/src/deepset_mcp/mcp/store.py similarity index 92% rename from src/deepset_mcp/store.py rename to src/deepset_mcp/mcp/store.py index 1f6ac5f..1974c19 100644 --- a/src/deepset_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/tool_factory.py b/src/deepset_mcp/mcp/tool_factory.py similarity index 90% rename from src/deepset_mcp/tool_factory.py rename to src/deepset_mcp/mcp/tool_factory.py index e321fe2..dbb0d71 100644 --- a/src/deepset_mcp/tool_factory.py +++ b/src/deepset_mcp/mcp/tool_factory.py @@ -14,6 +14,8 @@ 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, @@ -21,8 +23,6 @@ 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/tool_models.py b/src/deepset_mcp/mcp/tool_models.py similarity index 100% rename from src/deepset_mcp/tool_models.py rename to src/deepset_mcp/mcp/tool_models.py diff --git a/src/deepset_mcp/tool_registry.py b/src/deepset_mcp/mcp/tool_registry.py similarity index 99% rename from src/deepset_mcp/tool_registry.py rename to src/deepset_mcp/mcp/tool_registry.py index 846609e..11c6c66 100644 --- a/src/deepset_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 index 76a9d74..29972ad 100644 --- 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 index 5e4df42..7cac5fb 100644 --- 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,7 +15,7 @@ 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) @@ -23,7 +23,7 @@ def test_initialize_store_memory_backend(self) -> None: 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,10 +85,10 @@ 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 @@ -96,10 +96,10 @@ def test_initialize_store_caching(self) -> None: 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 index 32e91f4..b950294 100644 --- a/test/unit/test_tool_factory.py +++ b/test/unit/test_tool_factory.py @@ -11,15 +11,15 @@ import pytest from deepset_mcp.api.protocols import AsyncClientProtocol -from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore -from deepset_mcp.tool_factory import ( +from deepset_mcp.mcp.tool_factory import ( apply_client, apply_custom_args, apply_memory, apply_workspace, build_tool, ) -from deepset_mcp.tool_models import MemoryType, ToolConfig, WorkspaceMode +from deepset_mcp.mcp.tool_models import MemoryType, ToolConfig, WorkspaceMode +from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore from test.unit.conftest import BaseFakeClient @@ -89,7 +89,7 @@ async def sample_func(a: int) -> str: return str(a) config = ToolConfig(needs_workspace=False) - result = apply_workspace(sample_func, config, WorkspaceMode.STATIC) + result = apply_workspace(sample_func, config, workspace="some_workspace") assert result is sample_func @@ -100,7 +100,7 @@ async def sample_func(workspace: str, a: int) -> str: return f"{workspace}:{a}" config = ToolConfig(needs_workspace=True) - result = apply_workspace(sample_func, config, WorkspaceMode.DYNAMIC) + result = apply_workspace(sample_func, config) assert result is sample_func @@ -116,7 +116,7 @@ async def sample_func(workspace: str, a: int) -> str: return f"{workspace}:{a}" config = ToolConfig(needs_workspace=True) - result = apply_workspace(sample_func, config, WorkspaceMode.STATIC, "test-workspace") + result = apply_workspace(sample_func, config, "test-workspace") # Check signature was updated sig = inspect.signature(result) @@ -137,7 +137,7 @@ async def sample_func(workspace: str, a: int) -> str: return f"{workspace}:{a}" config = ToolConfig(needs_workspace=True) - result = apply_workspace(sample_func, config, WorkspaceMode.STATIC, "test-workspace") + result = apply_workspace(sample_func, config, "test-workspace") # Call should work with workspace injected output = await result(a=42) @@ -174,7 +174,7 @@ async def sample_func(a: int) -> str: with pytest.raises(ValueError, match="Invalid memory type"): apply_memory(sample_func, config, store) - @patch("deepset_mcp.tool_factory.explorable") + @patch("deepset_mcp.mcp.tool_factory.explorable") def test_explorable_memory_applied(self, mock_explorable: Any, store: ObjectStore) -> None: """Test that explorable decorator is applied.""" @@ -191,7 +191,7 @@ async def sample_func(a: int) -> str: mock_explorable.assert_called_once() mock_decorator.assert_called_once_with(sample_func) - @patch("deepset_mcp.tool_factory.referenceable") + @patch("deepset_mcp.mcp.tool_factory.referenceable") def test_referenceable_memory_applied(self, mock_referenceable: Any, store: ObjectStore) -> None: """Test that referenceable decorator is applied.""" @@ -208,7 +208,7 @@ async def sample_func(a: int) -> str: mock_referenceable.assert_called_once() mock_decorator.assert_called_once_with(sample_func) - @patch("deepset_mcp.tool_factory.explorable_and_referenceable") + @patch("deepset_mcp.mcp.tool_factory.explorable_and_referenceable") def test_both_memory_applied(self, mock_both: Any, store: ObjectStore) -> None: """Test that both memory decorator is applied.""" @@ -350,7 +350,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient to return our FakeClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42, ctx=mock_ctx) @@ -376,7 +376,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient to return our FakeClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42, ctx=mock_ctx) @@ -399,7 +399,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient to return our FakeClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42) @@ -426,7 +426,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient to return our FakeClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42, ctx=mock_ctx) @@ -451,7 +451,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient to return our FakeClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42) @@ -474,7 +474,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient to return our FakeClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42) @@ -588,7 +588,7 @@ async def sample_func(client: AsyncClientProtocol, workspace: str, a: int) -> st fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client output = await result(a=42, ctx=mock_ctx) @@ -628,7 +628,7 @@ async def sample_func(a: int) -> str: memory_type=MemoryType.EXPLORABLE, ) - with patch("deepset_mcp.tool_factory.explorable") as mock_explorable: + with patch("deepset_mcp.mcp.tool_factory.explorable") as mock_explorable: mock_decorator = MagicMock() mock_explorable.return_value = mock_decorator mock_decorator.return_value = sample_func @@ -684,7 +684,7 @@ async def sample_func(client: AsyncClientProtocol, a: int) -> str: # Mock the AsyncDeepsetClient fake_client = BaseFakeClient() - with patch("deepset_mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: + with patch("deepset_mcp.mcp.tool_factory.AsyncDeepsetClient") as mock_client_class: mock_client_class.return_value.__aenter__.return_value = fake_client await result(a=42, ctx=mock_ctx)