chore: rearrange some modules for better use

.gitignore

@@ -18,6 +18,7 @@ wheels/
 
 
 agent_runs/
+site/
 
 misc.xml
 modules.xml

pyproject.toml

@@ -141,5 +141,5 @@ ignore_missing_imports = true
 
 [[tool.mypy.overrides]]
 # These tests intentionally test runtime decorator behavior that mypy cannot understand
-module = ["test.unit.tools.tokonomics.*"]
+module = ["test.unit.tokonomics.*"]
 disable_error_code = ["arg-type", "attr-defined", "operator", "misc", "type-arg", "no-any-return", "assignment"]

src/deepset_mcp/store.py

@@ -7,7 +7,7 @@
 import functools
 import logging
 
-from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore, ObjectStoreBackend, RedisBackend
+from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore, ObjectStoreBackend, RedisBackend
 
 logger = logging.getLogger(__name__)
 

src/deepset_mcp/tools/tokonomics/__init__.py → src/deepset_mcp/tokonomics/__init__.py

@@ -4,12 +4,14 @@
 
 from .decorators import explorable, explorable_and_referenceable, referenceable
 from .explorer import RichExplorer
-from .object_store import InMemoryBackend, ObjectStore
+from .object_store import InMemoryBackend, ObjectStore, ObjectStoreBackend, RedisBackend
 
 __all__ = [
     # Core classes
     "InMemoryBackend",
+    "RedisBackend",
     "ObjectStore",
+    "ObjectStoreBackend",
     "RichExplorer",
     # Decorators
     "explorable",

src/deepset_mcp/tools/tokonomics/decorators.py → src/deepset_mcp/tokonomics/decorators.py

@@ -19,8 +19,8 @@
 
 from glom import GlomError, glom
 
-from deepset_mcp.tools.tokonomics.explorer import RichExplorer
-from deepset_mcp.tools.tokonomics.object_store import ObjectStore
+from deepset_mcp.tokonomics.explorer import RichExplorer
+from deepset_mcp.tokonomics.object_store import ObjectStore
 
 F = TypeVar("F", bound=Callable[..., Any])
 

src/deepset_mcp/tools/tokonomics/explorer.py → src/deepset_mcp/tokonomics/explorer.py

@@ -18,7 +18,7 @@
 from rich.console import Console
 from rich.pretty import Pretty
 
-from deepset_mcp.tools.tokonomics.object_store import ObjectStore
+from deepset_mcp.tokonomics.object_store import ObjectStore
 
 
 class RichExplorer:

src/deepset_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.tool_models import DeepsetDocsConfig, MemoryType, ToolConfig, WorkspaceMode
-from deepset_mcp.tool_registry import TOOL_REGISTRY
-from deepset_mcp.tools.tokonomics import (
+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]:
@@ -139,7 +139,7 @@ def apply_memory(
         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:
+    elif config.memory_type == MemoryType.EXPLORABLE_AND_REFERENCEABLE:
         return explorable_and_referenceable(object_store=store, explorer=explorer)(base_func)
     else:
         raise ValueError(f"Invalid memory type: {config.memory_type}")

src/deepset_mcp/tool_models.py

@@ -11,19 +11,41 @@ class MemoryType(StrEnum):
     """Configuration for how memory is provided to tools."""
 
     EXPLORABLE = "explorable"
+    """The tool's output is stored in the object store and nested properties can be explored through using the object
+    store tools.
+    """
+
     REFERENCEABLE = "referenceable"
-    BOTH = "both"
+    """The tool can be called by referencing an object or object-property that was stored in the object store."""
+
+    EXPLORABLE_AND_REFERENCEABLE = "explorable_and_referenceable"
+    """The tool's output is stored in the object store and it can be called by reference."""
+
     NO_MEMORY = "no_memory"
+    """The tool returns all outputs as is. It does not interact with the object store."""
 
 
 @dataclass
 class ToolConfig:
-    """Configuration for tool registration."""
+    """Configuration for tool registration.
+
+    It allows users to define what arguments should be passed to the tool at registration time. These arguments will not
+    be provided by the LLM as the tool will receive them programmatically through partial application.
+
+    The configuration also determines if a tool should store outputs in the object store.
+    """
 
     needs_client: bool = False
+    """If the tool should receive a configured instance of the 'AsyncDeepsetClient' at tool-registration time."""
+
     needs_workspace: bool = False
+    """If the tool should receive a static deepset workspace at tool-registration time."""
+
     memory_type: MemoryType = MemoryType.NO_MEMORY
+    """The type of memory this tool should use."""
+
     custom_args: dict[str, Any] = field(default_factory=dict)
+    """Any other arguments that should be passed to the tool at registration time instead of being passed by the LLM."""
 
 
 class WorkspaceMode(StrEnum):

src/deepset_mcp/tool_registry.py

@@ -89,7 +89,7 @@ async def search_docs(query: str) -> str:
         ToolConfig(
             needs_client=True,
             needs_workspace=True,
-            memory_type=MemoryType.BOTH,
+            memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE,
             custom_args={"skip_validation_errors": True},
         ),
     ),
@@ -98,7 +98,7 @@ async def search_docs(query: str) -> str:
         ToolConfig(
             needs_client=True,
             needs_workspace=True,
-            memory_type=MemoryType.BOTH,
+            memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE,
             custom_args={"skip_validation_errors": True},
         ),
     ),
@@ -117,7 +117,7 @@ async def search_docs(query: str) -> str:
     ),
     "validate_pipeline": (
         validate_pipeline_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE),
     ),
     "get_pipeline_logs": (
         get_pipeline_logs_tool,
@@ -137,11 +137,11 @@ async def search_docs(query: str) -> str:
     ),
     "create_index": (
         create_index_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE),
     ),
     "update_index": (
         update_index_tool,
-        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE),
     ),
     "deploy_index": (
         deploy_index_tool,
@@ -198,7 +198,7 @@ async def search_docs(query: str) -> str:
     ),
     "run_component": (
         run_component_tool,
-        ToolConfig(needs_client=True, memory_type=MemoryType.BOTH),
+        ToolConfig(needs_client=True, memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE),
     ),
     "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)),

src/deepset_mcp/tools/object_store.py

@@ -5,7 +5,7 @@
 from collections.abc import Callable
 from typing import Any
 
-from deepset_mcp.tools.tokonomics import RichExplorer
+from deepset_mcp.tokonomics import RichExplorer
 
 
 def create_get_from_object_store(explorer: RichExplorer) -> Callable[..., Any]:

test/integration/tools/tokonomics/test_integration_tokonomics.py → test/integration/tokonomics/test_integration_tokonomics.py

@@ -19,9 +19,8 @@
     docker = None
     NotFound = Exception
 
-from deepset_mcp.tools.tokonomics.decorators import explorable, explorable_and_referenceable, referenceable
-from deepset_mcp.tools.tokonomics.explorer import RichExplorer
-from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore, RedisBackend
+from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore, RedisBackend, RichExplorer
+from deepset_mcp.tokonomics.decorators import explorable, explorable_and_referenceable, referenceable
 
 pytestmark = pytest.mark.integration
 
@@ -65,7 +64,7 @@ def redis_container(docker_client: Any) -> Generator[Any, None, None]:
     for i in range(max_retries):
         try:
             # Try to connect
-            from deepset_mcp.tools.tokonomics.object_store import RedisBackend
+            from deepset_mcp.tokonomics import RedisBackend
 
             RedisBackend("redis://localhost:16379/0")
             break

test/unit/test_store.py

@@ -7,7 +7,7 @@
 import pytest
 
 from deepset_mcp.store import create_redis_backend, initialize_store
-from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore
+from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore
 
 
 class TestStoreInitialization:

test/unit/test_tool_factory.py

@@ -11,6 +11,7 @@
 import pytest
 
 from deepset_mcp.api.protocols import AsyncClientProtocol
+from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore
 from deepset_mcp.tool_factory import (
     apply_client,
     apply_custom_args,
@@ -19,7 +20,6 @@
     build_tool,
 )
 from deepset_mcp.tool_models import MemoryType, ToolConfig, WorkspaceMode
-from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore
 from test.unit.conftest import BaseFakeClient
 
 
@@ -219,7 +219,7 @@ async def sample_func(a: int) -> str:
         mock_both.return_value = mock_decorator
         mock_decorator.return_value = sample_func
 
-        config = ToolConfig(memory_type=MemoryType.BOTH)
+        config = ToolConfig(memory_type=MemoryType.EXPLORABLE_AND_REFERENCEABLE)
         apply_memory(sample_func, config, store)
 
         mock_both.assert_called_once()

test/unit/tools/tokonomics/test_decorators.py → test/unit/tokonomics/test_decorators.py

@@ -8,7 +8,8 @@
 
 import pytest
 
-from deepset_mcp.tools.tokonomics.decorators import (
+from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore, RichExplorer
+from deepset_mcp.tokonomics.decorators import (
     _add_str_to_type,
     _enhance_docstring_for_explorable,
     _enhance_docstring_for_references,
@@ -18,8 +19,6 @@
     explorable_and_referenceable,
     referenceable,
 )
-from deepset_mcp.tools.tokonomics.explorer import RichExplorer
-from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore
 
 
 class TestHelperFunctions:

test/unit/tools/tokonomics/test_explorer.py → test/unit/tokonomics/test_explorer.py

@@ -5,8 +5,7 @@
 import pytest
 from glom import Path
 
-from deepset_mcp.tools.tokonomics.explorer import RichExplorer
-from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore
+from deepset_mcp.tokonomics import InMemoryBackend, ObjectStore, RichExplorer
 
 
 class TestRichExplorer:

test/unit/tools/tokonomics/test_integration.py → test/unit/tokonomics/test_integration.py

@@ -8,14 +8,14 @@
 
 import pytest
 
-from deepset_mcp.tools.tokonomics import (
+from deepset_mcp.tokonomics import (
     InMemoryBackend,
     ObjectStore,
     RichExplorer,
     explorable,
     referenceable,
 )
-from deepset_mcp.tools.tokonomics.decorators import explorable_and_referenceable
+from deepset_mcp.tokonomics.decorators import explorable_and_referenceable
 
 
 class TestTokonomicsIntegration:

test/unit/tools/tokonomics/test_object_store.py → test/unit/tokonomics/test_object_store.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from deepset_mcp.tools.tokonomics.object_store import (
+from deepset_mcp.tokonomics import (
     InMemoryBackend,
     ObjectStore,
     RedisBackend,

test/unit/tools/tokonomics/test_object_store_backends.py → test/unit/tokonomics/test_object_store_backends.py

@@ -7,7 +7,7 @@
 
 import pytest
 
-from deepset_mcp.tools.tokonomics.object_store import (
+from deepset_mcp.tokonomics import (
     InMemoryBackend,
     RedisBackend,
 )