deepset-mcp 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

deepset_mcp/__init__.py

@@ -2,9 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import importlib.metadata
+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
 
-try:
-    __version__ = importlib.metadata.version(__name__)
-except importlib.metadata.PackageNotFoundError:
-    __version__ = "0.0.0"  # Fallback for development mode
+__all__ = ["configure_mcp_server", "WorkspaceMode", "ALL_DEEPSET_TOOLS", "DEEPSET_DOCS_DEFAULT_SHARE_URL"]

deepset_mcp/api/transport.py

@@ -6,6 +6,7 @@ import json
 import time
 from collections.abc import AsyncIterator
 from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Generic, Literal, Protocol, TypeVar, cast, overload
 
@@ -198,8 +199,10 @@ class AsyncTransport:
         config : dict, optional
             Configuration for httpx.AsyncClient, e.g., {'timeout': 10.0}
         """
-        config = config or {}
-        # Ensure auth header
+        # We deepcopy the config so that we don't mutate it when used for subsequent initializations
+        config = deepcopy(config) or {}
+
+        # Merge auth and other config headers
         headers = config.pop("headers", {})
         headers.setdefault("Authorization", f"Bearer {api_key}")
         # Build client kwargs

deepset_mcp/config.py

@@ -4,7 +4,14 @@
 
 """This module contains static configuration for the deepset MCP server."""
 
-from deepset_mcp import __version__
+import importlib.metadata
+
+try:
+    __version__ = importlib.metadata.version("deepset-mcp")
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"  # Fallback for development mode
+
+PACKAGE_VERSION = __version__
 
 # We need this mapping to which environment variables integrations are mapped to
 # The mapping is maintained in the pipeline operator:
@@ -29,5 +36,6 @@ TOKEN_DOMAIN_MAPPING = {
 }
 
 DEEPSET_DOCS_DEFAULT_SHARE_URL = "https://cloud.deepset.ai/shared_prototypes?share_token=prototype_eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3ODM0MjE0OTguNTk5LCJhdWQiOiJleHRlcm5hbCB1c2VyIiwiaXNzIjoiZEMiLCJ3b3Jrc3BhY2VfaWQiOiI4YzI0ZjExMi1iMjljLTQ5MWMtOTkzOS1hZTkxMDRhNTQyMWMiLCJ3b3Jrc3BhY2VfbmFtZSI6ImRjLWRvY3MtY29udGVudCIsIm9yZ2FuaXphdGlvbl9pZCI6ImNhOWYxNGQ0LWMyYzktNDYwZC04ZDI2LWY4Y2IwYWNhMDI0ZiIsInNoYXJlX2lkIjoiY2Y3MTA3ODAtOThmNi00MzlmLThiNzYtMmMwNDkyODNiMDZhIiwibG9naW5fcmVxdWlyZWQiOmZhbHNlfQ.5j6DCNRQ1_KB8lhIJqHyw2hBIleEW1_Y_UBuH6MTYY0"
+DOCS_SEARCH_TOOL_NAME = "search_docs"
 
 DEFAULT_CLIENT_HEADER = {"headers": {"User-Agent": f"deepset-mcp/{__version__}"}}

deepset_mcp/main.py

@@ -2,191 +2,208 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import argparse
-import asyncio
 import logging
 import os
-from pathlib import Path
-from urllib.parse import parse_qs, urlparse
+from enum import StrEnum
+from typing import Annotated
 
-import jwt
+import typer
 from mcp.server.fastmcp import FastMCP
 
-from deepset_mcp.api.client import AsyncDeepsetClient
-from deepset_mcp.config import DEEPSET_DOCS_DEFAULT_SHARE_URL
-from deepset_mcp.tool_factory import WorkspaceMode, register_tools
+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
 
-# Initialize MCP Server
-mcp = FastMCP("Deepset Cloud MCP", settings={"log_level": "ERROR"})
 
-logging.getLogger("uvicorn").setLevel(logging.WARNING)
-logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
-logging.getLogger("fastapi").setLevel(logging.WARNING)
-logging.getLogger("httpx").setLevel(logging.WARNING)
-logging.getLogger("httpcore").setLevel(logging.WARNING)
-logging.getLogger("mcp").setLevel(logging.WARNING)
+class TransportEnum(StrEnum):
+    """Transport mode for the MCP server."""
 
+    STDIO = "stdio"
+    STREAMABLE_HTTP = "streamable-http"
 
-@mcp.prompt()
-async def deepset_copilot() -> str:
-    """System prompt for the deepset copilot."""
-    prompt_path = Path(__file__).parent / "prompts/deepset_copilot_prompt.md"
 
-    return prompt_path.read_text()
+app = typer.Typer(
+    name="deepset-mcp",
+    help="Run the Deepset MCP server to interact with the deepset AI platform.",
+    no_args_is_help=True,
+)
 
 
-@mcp.prompt()
-async def deepset_recommended_prompt() -> str:
-    """Recommended system prompt for the deepset copilot."""
-    prompt_path = Path(__file__).parent / "prompts/deepset_debugging_agent.md"
-
-    return prompt_path.read_text()
-
-
-async def fetch_shared_prototype_details(share_url: str) -> tuple[str, str, str]:
-    """Gets the pipeline name, workspace name and an API token for a shared prototype url.
-
-    :param share_url: The URL of a shared prototype on the deepset platform.
-
-    :returns: A tuple containing the pipeline name, workspace name and an API token.
-    """
-    parsed_url = urlparse(share_url)
-    query_params = parse_qs(parsed_url.query)
-    share_token = query_params.get("share_token", [None])[0]
-    if not share_token:
-        raise ValueError("Invalid share URL: missing share_token parameter.")
-
-    jwt_token = share_token.replace("prototype_", "")
-
-    decoded_token = jwt.decode(jwt_token, options={"verify_signature": False})
-    workspace_name = decoded_token.get("workspace_name")
-    if not workspace_name:
-        raise ValueError("Invalid JWT in share_token: missing 'workspace_name'.")
-
-    share_id = decoded_token.get("share_id")
-    if not share_id:
-        raise ValueError("Invalid JWT in share_token: missing 'share_id'.")
-
-    # For shared prototypes, we need to:
-    # 1. Fetch prototype details (pipeline name) using the information encoded in the JWT
-    # 2. Create a shared prototype user
-    async with AsyncDeepsetClient(api_key=share_token) as client:
-        response = await client.request(f"/v1/workspaces/{workspace_name}/shared_prototypes/{share_id}")
-        if not response.success:
-            raise ValueError(f"Failed to fetch shared prototype details: {response.status_code} {response.json}")
-
-        data = response.json or {}
-        pipeline_names: list[str] = data.get("pipeline_names", [])
-        if not pipeline_names:
-            raise ValueError("No pipeline names found in shared prototype response.")
-
-        user_info = await client.request("/v1/workspaces/dc-docs-content/shared_prototype_users", method="POST")
-
-        if not user_info.success:
-            raise ValueError("Failed to fetch user information from shared prototype response.")
-
-        user_data = user_info.json or {}
-
-        try:
-            api_key = user_data["user_token"]
-        except KeyError:
-            raise ValueError("No user token in shared prototype response.") from None
-
-    return workspace_name, pipeline_names[0], api_key
-
-
-def main() -> None:
-    """Entrypoint for the deepset MCP server."""
-    parser = argparse.ArgumentParser(description="Run the Deepset MCP server.")
-    parser.add_argument(
-        "--workspace",
-        "-w",
-        help="Deepset workspace (env DEEPSET_WORKSPACE)",
-    )
-    parser.add_argument(
-        "--api-key",
-        "-k",
-        help="Deepset API key (env DEEPSET_API_KEY)",
-    )
-    parser.add_argument(
-        "--docs-share-url",
-        "-d",
-        default=DEEPSET_DOCS_DEFAULT_SHARE_URL,
-        help="Deepset docs search share URL (env DEEPSET_DOCS_SHARE_URL)",
-    )
-    parser.add_argument(
-        "--workspace-mode",
-        "-m",
-        choices=[WorkspaceMode.STATIC, WorkspaceMode.DYNAMIC],
-        default=WorkspaceMode.STATIC,
-        help=(
-            "Whether workspace should be set statically or dynamically provided during a tool call. "
-            f"Default: '{WorkspaceMode.STATIC}'"
+@app.command()
+def main(
+    workspace: Annotated[
+        str | None,
+        typer.Option(
+            "--workspace",
+            help="Deepset workspace name. Can also be set via DEEPSET_WORKSPACE environment variable.",
         ),
-    )
-    parser.add_argument(
-        "--tools",
-        "-t",
-        nargs="*",
-        help="Space-separated list of tools to register (default: all)",
-    )
-    parser.add_argument(
-        "--list-tools",
-        "-l",
-        action="store_true",
-        help="List all available tools and exit",
-    )
-    args = parser.parse_args()
-
+    ] = None,
+    api_key: Annotated[
+        str | None,
+        typer.Option(
+            "--api-key",
+            help="Deepset API key for authentication. Can also be set via DEEPSET_API_KEY environment variable.",
+        ),
+    ] = None,
+    api_url: Annotated[
+        str | None,
+        typer.Option(
+            "--api-url",
+            help="Deepset API base URL. Can also be set via DEEPSET_API_URL environment variable.",
+        ),
+    ] = None,
+    docs_share_url: Annotated[
+        str | None,
+        typer.Option(
+            "--docs-share-url",
+            help="Deepset docs search share URL. Can also be set via DEEPSET_DOCS_SHARE_URL environment variable.",
+        ),
+    ] = None,
+    workspace_mode: Annotated[
+        WorkspaceMode,
+        typer.Option(
+            "--workspace-mode",
+            help="Whether workspace should be set statically or dynamically provided during a tool call.",
+        ),
+    ] = WorkspaceMode.STATIC,
+    tools: Annotated[
+        list[str] | None,
+        typer.Option(
+            "--tools",
+            help="Space-separated list of tools to register. If not specified, all available tools will be registered.",
+        ),
+    ] = None,
+    list_tools: Annotated[
+        bool,
+        typer.Option(
+            "--list-tools",
+            help="List all available tools and exit.",
+        ),
+    ] = False,
+    api_key_from_auth_header: Annotated[
+        bool,
+        typer.Option(
+            "--api-key-from-auth-header/--no-api-key-from-auth-header",
+            help="Get the deepset API key from the request's authorization header instead of using a static key.",
+        ),
+    ] = False,
+    transport: Annotated[
+        TransportEnum,
+        typer.Option(
+            "--transport",
+            help="The type of transport to use for running the MCP server.",
+        ),
+    ] = TransportEnum.STDIO,
+    object_store_backend: Annotated[
+        str | None,
+        typer.Option(
+            "--object-store-backend",
+            help="Object store backend type: 'memory' or 'redis'. "
+            "Can also be set via OBJECT_STORE_BACKEND environment variable. Default is 'memory'.",
+        ),
+    ] = None,
+    object_store_redis_url: Annotated[
+        str | None,
+        typer.Option(
+            "--object-store-redis-url",
+            help="Redis connection URL (e.g., redis://localhost:6379). "
+            "Can also be set via OBJECT_STORE_REDIS_URL environment variable.",
+        ),
+    ] = None,
+    object_store_ttl: Annotated[
+        int,
+        typer.Option(
+            "--object-store-ttl",
+            help="TTL in seconds for stored objects. Default: 600 (10 minutes). "
+            "Can also be set via OBJECT_STORE_TTL environment variable.",
+        ),
+    ] = 600,
+) -> None:
+    """
+    Run the Deepset MCP server.
+
+    The Deepset MCP server provides tools to interact with the deepset AI platform,
+    allowing you to create, debug, and learn about pipelines on the platform.
+
+    :param workspace: Deepset workspace name
+    :param api_key: Deepset API key for authentication
+    :param api_url: Deepset API base URL
+    :param docs_share_url: Deepset docs search share URL
+    :param workspace_mode: Whether workspace should be set statically or dynamically
+    :param tools: List of tools to register
+    :param list_tools: List all available tools and exit
+    :param api_key_from_auth_header: Get API key from authorization header
+    :param transport: Type of transport to use for the MCP server
+    :param object_store_backend: Object store backend type ('memory' or 'redis')
+    :param object_store_redis_url: Redis connection URL (required if backend='redis')
+    :param object_store_ttl: TTL in seconds for stored objects
+    """
     # Handle --list-tools flag early
-    if args.list_tools:
-        from deepset_mcp.tool_factory import TOOL_REGISTRY
-
-        print("Available tools:")
+    if list_tools:
+        typer.echo("Available tools:")
         for tool_name in sorted(TOOL_REGISTRY.keys()):
-            print(f"  {tool_name}")
-        return
-
-    # prefer flags, fallback to env
-    workspace = args.workspace or os.getenv("DEEPSET_WORKSPACE")
-    api_key = args.api_key or os.getenv("DEEPSET_API_KEY")
-    docs_share_url = args.docs_share_url or os.getenv("DEEPSET_DOCS_SHARE_URL")
-
-    if docs_share_url:
-        try:
-            workspace_name, pipeline_name, api_key_docs = asyncio.run(fetch_shared_prototype_details(docs_share_url))
-            os.environ["DEEPSET_DOCS_WORKSPACE"] = workspace_name
-            os.environ["DEEPSET_DOCS_PIPELINE_NAME"] = pipeline_name
-            os.environ["DEEPSET_DOCS_API_KEY"] = api_key_docs
-        except (ValueError, jwt.DecodeError) as e:
-            parser.error(f"Error processing --docs-share-url: {e}")
-
-    # Create server configuration
-    workspace_mode = WorkspaceMode(args.workspace_mode)
-
-    if workspace_mode == WorkspaceMode.STATIC:
-        if not workspace:
-            parser.error("Missing workspace: set --workspace or DEEPSET_WORKSPACE")
-
-    if not api_key:
-        parser.error("Missing API key: set --api-key or DEEPSET_API_KEY")
-
-    # make sure downstream tools see them (for implicit mode)
-    if workspace:
-        os.environ["DEEPSET_WORKSPACE"] = workspace
-    os.environ["DEEPSET_API_KEY"] = api_key
-
-    # Parse tool names if provided
-    tool_names = None
-    if args.tools:
-        tool_names = set(args.tools)
-
-    # Register tools based on configuration
-    register_tools(mcp, workspace_mode, workspace, tool_names)
+            typer.echo(f"  {tool_name}")
+        raise typer.Exit()
+
+    # Prefer command line arguments, fallback to environment variables
+    workspace = workspace or os.getenv("DEEPSET_WORKSPACE")
+    api_key = api_key or os.getenv("DEEPSET_API_KEY")
+    api_url = api_url or os.getenv("DEEPSET_API_URL")
+    docs_share_url = docs_share_url or os.getenv("DEEPSET_DOCS_SHARE_URL", DEEPSET_DOCS_DEFAULT_SHARE_URL)
+
+    # ObjectStore configuration
+    backend = str(object_store_backend or os.getenv("OBJECT_STORE_BACKEND", "memory"))
+    redis_url = object_store_redis_url or os.getenv("OBJECT_STORE_REDIS_URL")
+    ttl = int(os.getenv("OBJECT_STORE_TTL", str(object_store_ttl)))
+
+    if tools:
+        tool_names = set(tools)
+    else:
+        logging.info("Registering all available tools.")
+        tool_names = set(TOOL_REGISTRY.keys())
+
+    if api_key is None and not api_key_from_auth_header:
+        typer.echo(
+            "Error: API key is required. Either provide --api-key or use --api-key-from-auth-header "
+            "to fetch the API key from the authorization header.",
+            err=True,
+        )
+        raise typer.Exit(1)
+
+    if workspace_mode == WorkspaceMode.STATIC and not workspace:
+        typer.echo(
+            "Error: Workspace is required when using static workspace mode. "
+            "Set --workspace or DEEPSET_WORKSPACE environment variable.",
+            err=True,
+        )
+        raise typer.Exit(1)
+
+    if DOCS_SEARCH_TOOL_NAME in tool_names and docs_share_url is None:
+        typer.echo(
+            f"Error: {DOCS_SEARCH_TOOL_NAME} tool is requested but no docs share URL provided. "
+            "Set --docs-share-url or DEEPSET_DOCS_SHARE_URL environment variable.",
+            err=True,
+        )
+        raise typer.Exit(1)
+
+    mcp = FastMCP("deepset AI platform MCP server")
+    configure_mcp_server(
+        mcp_server_instance=mcp,
+        workspace_mode=workspace_mode,
+        deepset_api_key=api_key,
+        deepset_api_url=api_url,
+        deepset_workspace=workspace,
+        tools_to_register=tool_names,
+        deepset_docs_shareable_prototype_url=docs_share_url,
+        get_api_key_from_authorization_header=api_key_from_auth_header,
+        object_store_backend=backend,
+        object_store_redis_url=redis_url,
+        object_store_ttl=ttl,
+    )
 
-    # run with SSE transport (HTTP+Server-Sent Events)
-    mcp.run(transport="stdio")
+    mcp.run(transport=transport.value)
 
 
 if __name__ == "__main__":
-    main()
+    app()

deepset_mcp/server.py

@@ -0,0 +1,154 @@
+# SPDX-FileCopyrightText: 2025-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import asyncio
+from urllib.parse import parse_qs, urlparse
+
+import jwt
+from mcp.server.fastmcp import FastMCP
+
+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
+
+
+def configure_mcp_server(
+    mcp_server_instance: FastMCP,
+    workspace_mode: WorkspaceMode | str,
+    tools_to_register: set[str] | None = None,
+    deepset_api_key: str | None = None,
+    deepset_api_url: str | None = None,
+    deepset_workspace: str | None = None,
+    deepset_docs_shareable_prototype_url: str | None = None,
+    get_api_key_from_authorization_header: bool = False,
+    object_store_backend: str = "memory",
+    object_store_redis_url: str | None = None,
+    object_store_ttl: int = 600,
+) -> None:
+    """Configure the MCP server with the specified tools and settings.
+
+    :param mcp_server_instance: The FastMCP server instance to configure
+    :param workspace_mode: The workspace mode ("static" or "dynamic")
+    :param tools_to_register: Set of tool names to register with the server.
+        Will register all tools if set to None.
+    :param deepset_api_key: Optional Deepset API key for authentication
+    :param deepset_api_url: Optional Deepset API base URL
+    :param deepset_workspace: Optional workspace name for static mode
+    :param deepset_docs_shareable_prototype_url: Shareable prototype URL that allows access to a docs search pipeline.
+        Will fall back to the default shareable prototype URL if set to None.
+    :param get_api_key_from_authorization_header: Whether to extract API key from authorization header
+    :param object_store_backend: Object store backend type ('memory' or 'redis')
+    :param object_store_redis_url: Redis connection URL (required if backend='redis')
+    :param object_store_ttl: TTL in seconds for stored objects
+    :raises ValueError: If required parameters are missing or invalid
+    """
+    if tools_to_register is None:
+        tools_to_register = set(TOOL_REGISTRY.keys())
+
+    if deepset_docs_shareable_prototype_url is None:
+        deepset_docs_shareable_prototype_url = DEEPSET_DOCS_DEFAULT_SHARE_URL
+
+    if isinstance(workspace_mode, str):
+        workspace_mode = WorkspaceMode(workspace_mode)
+
+    if workspace_mode == WorkspaceMode.STATIC and deepset_workspace is None:
+        raise ValueError(
+            "Static workspace mode requires a workspace name. "
+            "Please provide 'deepset_workspace' when using WorkspaceMode.STATIC."
+        )
+
+    if deepset_api_key is None and not get_api_key_from_authorization_header:
+        raise ValueError(
+            "API key is required for authentication. "
+            "Please provide 'deepset_api_key' or enable 'get_api_key_from_authorization_header'."
+        )
+
+    workspace_name, pipeline_name, api_key_docs = asyncio.run(
+        fetch_shared_prototype_details(deepset_docs_shareable_prototype_url)
+    )
+    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)
+
+    register_tools(
+        mcp_server_instance=mcp_server_instance,
+        workspace_mode=workspace_mode,
+        workspace=deepset_workspace,
+        tool_names=tools_to_register,
+        docs_config=docs_config,
+        get_api_key_from_authorization_header=get_api_key_from_authorization_header,
+        api_key=deepset_api_key,
+        base_url=deepset_api_url,
+        object_store=store,
+    )
+
+
+async def fetch_shared_prototype_details(share_url: str) -> tuple[str, str, str]:
+    """Extract pipeline name, workspace name and API token from a shared prototype URL.
+
+    :param share_url: The URL of a shared prototype on the Deepset platform
+    :returns: A tuple containing (workspace_name, pipeline_name, api_key)
+    :raises ValueError: If the URL is invalid or missing required parameters
+    """
+    parsed_url = urlparse(share_url)
+    query_params = parse_qs(parsed_url.query)
+    share_token = query_params.get("share_token", [None])[0]
+    if not share_token:
+        raise ValueError(
+            "Invalid share URL: missing 'share_token' parameter. Please provide a valid Deepset prototype share URL."
+        )
+
+    jwt_token = share_token.replace("prototype_", "")
+
+    decoded_token = jwt.decode(jwt_token, options={"verify_signature": False})
+    workspace_name = decoded_token.get("workspace_name")
+    if not workspace_name:
+        raise ValueError(
+            "Invalid share token: missing 'workspace_name' in JWT. The provided share URL may be corrupted or expired."
+        )
+
+    share_id = decoded_token.get("share_id")
+    if not share_id:
+        raise ValueError(
+            "Invalid share token: missing 'share_id' in JWT. The provided share URL may be corrupted or expired."
+        )
+
+    # For shared prototypes, we need to:
+    # 1. Fetch prototype details (pipeline name) using the information encoded in the JWT
+    # 2. Create a shared prototype user
+    async with AsyncDeepsetClient(api_key=share_token) as client:
+        response = await client.request(f"/v1/workspaces/{workspace_name}/shared_prototypes/{share_id}")
+        if not response.success:
+            raise ValueError(
+                f"Failed to fetch shared prototype details: HTTP {response.status_code}. Response: {response.json}"
+            )
+
+        data = response.json or {}
+        pipeline_names: list[str] = data.get("pipeline_names", [])
+        if not pipeline_names:
+            raise ValueError(
+                "No pipeline names found in shared prototype. The prototype may not be properly configured."
+            )
+
+        user_info = await client.request("/v1/workspaces/dc-docs-content/shared_prototype_users", method="POST")
+
+        if not user_info.success:
+            raise ValueError(
+                f"Failed to create user session for shared prototype. HTTP {user_info.status_code}: {user_info.json}"
+            )
+
+        user_data = user_info.json or {}
+
+        try:
+            api_key = user_data["user_token"]
+        except KeyError:
+            raise ValueError(
+                "No user token found in shared prototype response. Unable to authenticate with the prototype."
+            ) from None
+
+    return workspace_name, pipeline_names[0], api_key

deepset_mcp/store.py

@@ -4,6 +4,55 @@
 
 """Global store for the MCP server."""
 
-from deepset_mcp.tools.tokonomics import ObjectStore
+import functools
+import logging
 
-STORE = ObjectStore()
+from deepset_mcp.tools.tokonomics.object_store import InMemoryBackend, ObjectStore, ObjectStoreBackend, RedisBackend
+
+logger = logging.getLogger(__name__)
+
+
+def create_redis_backend(url: str) -> ObjectStoreBackend:
+    """Create Redis backend, failing if connection fails.
+
+    :param url: Redis connection URL
+    :raises Exception: If Redis connection fails
+    """
+    try:
+        backend = RedisBackend(url)
+        logger.info(f"Successfully connected to Redis at {url} (using UUIDs for IDs)")
+
+        return backend
+
+    except Exception as e:
+        logger.error(f"Failed to connect to Redis at {url}: {e}")
+        raise
+
+
+@functools.lru_cache(maxsize=1)
+def initialize_store(
+    backend: str = "memory",
+    redis_url: str | None = None,
+    ttl: int = 600,
+) -> ObjectStore:
+    """Initialize the object store.
+
+    :param backend: Backend type ('memory' or 'redis')
+    :param redis_url: Redis connection URL (required if backend='redis')
+    :param ttl: Time-to-live in seconds for stored objects
+    :raises ValueError: If Redis backend selected but no URL provided
+    :raises Exception: If Redis connection fails
+    """
+    if backend == "redis":
+        if not redis_url:
+            raise ValueError("'redis_url' is None. Provide a 'redis_url' to use the redis backend.")
+        backend_instance = create_redis_backend(redis_url)
+
+    else:
+        logger.info("Using in-memory backend")
+        backend_instance = InMemoryBackend()
+
+    store = ObjectStore(backend=backend_instance, ttl=ttl)
+    logger.info(f"Initialized ObjectStore with {backend} backend and TTL={ttl}s")
+
+    return store