Skip to content

Commit b1b8683

Browse files
committed
Reapply "feat: add redis backend"
This reverts commit 3691ec9.
1 parent 5e6be09 commit b1b8683

12 files changed

Lines changed: 567 additions & 751 deletions

src/deepset_mcp/implementation_plan.md

Lines changed: 0 additions & 516 deletions
This file was deleted.

src/deepset_mcp/main.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,7 @@ def main(
155155
# ObjectStore configuration
156156
backend = object_store_backend or os.getenv("DEEPSET_OBJECT_STORE_BACKEND", "memory")
157157
redis_url = redis_url or os.getenv("DEEPSET_REDIS_URL")
158-
ttl = float(os.getenv("DEEPSET_OBJECT_STORE_TTL", str(object_store_ttl)))
158+
ttl = int(os.getenv("DEEPSET_OBJECT_STORE_TTL", str(object_store_ttl)))
159159

160160
if tools:
161161
tool_names = set(tools)

src/deepset_mcp/tool_factory.py

Lines changed: 9 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66

77
import functools
88
import inspect
9-
import os
109
import re
1110
from collections.abc import Awaitable, Callable
1211
from typing import Any
@@ -26,23 +25,6 @@
2625
)
2726

2827

29-
def are_docs_available() -> bool:
30-
"""Checks if documentation search is available."""
31-
return bool(
32-
os.environ.get("DEEPSET_DOCS_WORKSPACE", False)
33-
and os.environ.get("DEEPSET_DOCS_PIPELINE_NAME", False)
34-
and os.environ.get("DEEPSET_DOCS_API_KEY", False)
35-
)
36-
37-
38-
def get_workspace_from_env() -> str:
39-
"""Gets the workspace configured from environment variable."""
40-
workspace = os.environ.get("DEEPSET_WORKSPACE")
41-
if not workspace:
42-
raise ValueError("DEEPSET_WORKSPACE environment variable not set")
43-
return workspace
44-
45-
4628
def apply_custom_args(base_func: Callable[..., Any], config: ToolConfig) -> Callable[..., Any]:
4729
"""
4830
Applies custom keyword arguments defined in the ToolConfig to a function.
@@ -117,8 +99,7 @@ def apply_workspace(
11799

118100
@functools.wraps(base_func)
119101
async def workspace_wrapper(*args: Any, **kwargs: Any) -> Any:
120-
ws = workspace or get_workspace_from_env()
121-
return await base_func(*args, workspace=ws, **kwargs)
102+
return await base_func(*args, workspace=workspace, **kwargs)
122103

123104
# Remove workspace from signature
124105
original_sig = inspect.signature(base_func)
@@ -165,7 +146,11 @@ def apply_memory(
165146

166147

167148
def apply_client(
168-
base_func: Callable[..., Any], config: ToolConfig, use_request_context: bool = True, base_url: str | None = None
149+
base_func: Callable[..., Any],
150+
config: ToolConfig,
151+
use_request_context: bool = True,
152+
base_url: str | None = None,
153+
api_key: str | None = None,
169154
) -> Callable[..., Any]:
170155
"""
171156
Applies the deepset API client to a function.
@@ -178,6 +163,7 @@ def apply_client(
178163
:param config: The ToolConfig for the function.
179164
:param use_request_context: Whether to collect the API key from the request context.
180165
:param base_url: Base URL for the deepset API.
166+
:param api_key: The API key to use.
181167
:returns: Function with client injection applied and updated signature/docstring.
182168
:raises ValueError: If API key cannot be extracted from request context.
183169
"""
@@ -213,26 +199,14 @@ async def client_wrapper_with_context(*args: Any, **kwargs: Any) -> Any:
213199
ctx_param = inspect.Parameter(name="ctx", kind=inspect.Parameter.KEYWORD_ONLY, annotation=Context)
214200
new_params.append(ctx_param)
215201
client_wrapper_with_context.__signature__ = original_sig.replace(parameters=new_params) # type: ignore
216-
217-
# Remove client from docstring
218-
if base_func.__doc__:
219-
import re
220-
221-
doc = base_func.__doc__
222-
doc = re.sub(
223-
r"^\s*:param\s+client.*?(?=^\s*:|^\s*$|\Z)",
224-
"",
225-
doc,
226-
flags=re.MULTILINE | re.DOTALL,
227-
)
228-
client_wrapper_with_context.__doc__ = "\n".join([line.rstrip() for line in doc.strip().split("\n")])
202+
client_wrapper_with_context.__doc__ = remove_params_from_docstring(base_func.__doc__, {"client"})
229203

230204
return client_wrapper_with_context
231205
else:
232206

233207
@functools.wraps(base_func)
234208
async def client_wrapper_without_context(*args: Any, **kwargs: Any) -> Any:
235-
client_kwargs: dict[str, Any] = {"transport_config": DEFAULT_CLIENT_HEADER}
209+
client_kwargs: dict[str, Any] = {"transport_config": DEFAULT_CLIENT_HEADER, "api_key": api_key}
236210
if base_url:
237211
client_kwargs["base_url"] = base_url
238212
async with AsyncDeepsetClient(**client_kwargs) as client:

src/deepset_mcp/tools/tokonomics/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,11 +60,12 @@
6060

6161
from .decorators import explorable, explorable_and_referenceable, referenceable
6262
from .explorer import RichExplorer
63-
from .object_store import Explorable, ObjectRef, ObjectStore
63+
from .object_store import Explorable, InMemoryBackend, ObjectRef, ObjectStore
6464

6565
__all__ = [
6666
# Core classes
6767
"Explorable",
68+
"InMemoryBackend",
6869
"ObjectRef",
6970
"ObjectStore",
7071
"RichExplorer",

src/deepset_mcp/tools/tokonomics/object_store.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ def __init__(self, redis_url: str) -> None:
154154
"Redis package not installed. Install with: pip install deepset-mcp[redis] to use the RedisBackend."
155155
) from e
156156

157-
self._client = redis.from_url(redis_url, decode_responses=False)
157+
self._client = redis.from_url(redis_url, decode_responses=False) # type: ignore[no-untyped-call]
158158
# Test connection immediately
159159
self._client.ping()
160160

@@ -172,7 +172,7 @@ def set(self, key: str, value: bytes, ttl_seconds: int | None) -> None:
172172

173173
def get(self, key: str) -> bytes | None:
174174
"""Get a value at key."""
175-
return self._client.get(key)
175+
return self._client.get(key) # type: ignore[no-any-return]
176176

177177
def delete(self, key: str) -> bool:
178178
"""Delete a value at key."""

test/unit/test_store.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
from unittest.mock import MagicMock, patch
6+
7+
import pytest
8+
9+
from deepset_mcp.store import create_redis_backend, initialize_store
10+
from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore
11+
12+
13+
class TestStoreInitialization:
14+
"""Test store initialization functions."""
15+
16+
def test_initialize_store_memory_backend(self) -> None:
17+
"""Test initializing store with memory backend."""
18+
store = initialize_store(backend="memory", ttl=1800)
19+
20+
assert isinstance(store, ObjectStore)
21+
assert isinstance(store._backend, InMemoryBackend)
22+
assert store._ttl == 1800
23+
24+
def test_initialize_store_default_backend(self) -> None:
25+
"""Test initializing store with default backend."""
26+
store = initialize_store()
27+
28+
assert isinstance(store, ObjectStore)
29+
assert isinstance(store._backend, InMemoryBackend)
30+
assert store._ttl == 600
31+
32+
def test_initialize_store_redis_backend_success(self) -> None:
33+
"""Test initializing store with Redis backend successfully."""
34+
mock_redis_client = MagicMock()
35+
mock_redis_client.ping.return_value = True
36+
37+
with patch("redis.from_url", return_value=mock_redis_client):
38+
store = initialize_store(backend="redis", redis_url="redis://localhost:6379", ttl=7200)
39+
40+
assert isinstance(store, ObjectStore)
41+
assert store._ttl == 7200
42+
mock_redis_client.ping.assert_called_once()
43+
44+
def test_initialize_store_redis_backend_no_url(self) -> None:
45+
"""Test initializing store with Redis backend but no URL."""
46+
with pytest.raises(ValueError, match="redis_url.*is None"):
47+
initialize_store(backend="redis", redis_url=None)
48+
49+
def test_initialize_store_redis_backend_connection_failure(self) -> None:
50+
"""Test initializing store with Redis backend connection failure."""
51+
mock_redis_client = MagicMock()
52+
mock_redis_client.ping.side_effect = Exception("Connection failed")
53+
54+
with patch("redis.from_url", return_value=mock_redis_client):
55+
with pytest.raises(Exception, match="Connection failed"):
56+
initialize_store(backend="redis", redis_url="redis://localhost:6379")
57+
58+
def test_create_redis_backend_success(self) -> None:
59+
"""Test creating Redis backend successfully."""
60+
mock_redis_client = MagicMock()
61+
mock_redis_client.ping.return_value = True
62+
63+
with patch("redis.from_url", return_value=mock_redis_client):
64+
backend = create_redis_backend("redis://localhost:6379")
65+
66+
assert backend is not None
67+
mock_redis_client.ping.assert_called_once()
68+
69+
def test_create_redis_backend_redis_not_installed(self) -> None:
70+
"""Test creating Redis backend when redis package is not installed."""
71+
with patch("builtins.__import__", side_effect=ImportError("No module named 'redis'")):
72+
with pytest.raises(ImportError, match="Redis package not installed"):
73+
create_redis_backend("redis://localhost:6379")
74+
75+
def test_create_redis_backend_connection_failure(self) -> None:
76+
"""Test creating Redis backend with connection failure."""
77+
mock_redis = MagicMock()
78+
mock_redis.from_url.return_value = mock_redis
79+
mock_redis.ping.side_effect = Exception("Connection failed")
80+
81+
with patch("redis.from_url", return_value=mock_redis):
82+
with pytest.raises(Exception, match="Connection failed"):
83+
create_redis_backend("redis://localhost:6379")
84+
85+
def test_initialize_store_caching(self) -> None:
86+
"""Test that initialize_store uses caching."""
87+
# Clear any existing cache
88+
initialize_store.cache_clear()
89+
90+
store1 = initialize_store(backend="memory", ttl=1800)
91+
store2 = initialize_store(backend="memory", ttl=1800)
92+
93+
# Should return the same instance due to caching
94+
assert store1 is store2
95+
96+
def test_initialize_store_different_params_different_instances(self) -> None:
97+
"""Test that different parameters create different instances."""
98+
# Clear any existing cache
99+
initialize_store.cache_clear()
100+
101+
store1 = initialize_store(backend="memory", ttl=1800)
102+
store2 = initialize_store(backend="memory", ttl=3600)
103+
104+
# Should return different instances due to different parameters
105+
assert store1 is not store2
106+
assert store1._ttl == 1800
107+
assert store2._ttl == 3600
108+
109+
def test_initialize_store_unknown_backend(self) -> None:
110+
"""Test initializing store with unknown backend defaults to memory."""
111+
store = initialize_store(backend="unknown", ttl=1800)
112+
113+
assert isinstance(store, ObjectStore)
114+
assert isinstance(store._backend, InMemoryBackend)
115+
assert store._ttl == 1800

0 commit comments

Comments
 (0)