Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/deepset_mcp/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
6 changes: 3 additions & 3 deletions src/deepset_mcp/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
10 changes: 10 additions & 0 deletions src/deepset_mcp/mcp/__init__.py
Original file line number Diff line number Diff line change
@@ -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"]
12 changes: 7 additions & 5 deletions src/deepset_mcp/server.py → src/deepset_mcp/mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/deepset_mcp/store.py → src/deepset_mcp/mcp/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions test/unit/test_server_base_url.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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()
Expand Down
26 changes: 13 additions & 13 deletions test/unit/test_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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."""
Expand All @@ -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."""
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down
Loading