feat: migrate index tools to decorator pattern (#111)

Author: mathislucka

src/deepset_mcp/tool_factory.py

@@ -5,7 +5,7 @@
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from enum import StrEnum
-from typing import Any, Literal
+from typing import Any
 
 from mcp.server.fastmcp import FastMCP
 
@@ -98,13 +98,22 @@ class WorkspaceMode(StrEnum):
     EXPLICIT = "explicit"  # workspace as required parameter in tool signature
 
 
+class MemoryType(StrEnum):
+    """Configuration for how memory is provided to tools."""
+
+    EXPLORABLE = "explorable"
+    REFERENCEABLE = "referenceable"
+    BOTH = "both"
+    NO_MEMORY = "no_memory"
+
+
 @dataclass
 class ToolConfig:
     """Configuration for tool registration."""
 
     needs_client: bool = False
     needs_workspace: bool = False
-    decorator_type: Literal["none", "explorable", "referenceable", "both"] = "none"
+    memory_type: MemoryType = MemoryType.NO_MEMORY
     custom_args: dict[str, Any] | None = None  # For special cases like search_component_definition
 
 
@@ -127,11 +136,26 @@ def get_workspace_from_env() -> str:
     "validate_pipeline": (validate_pipeline_tool, ToolConfig(needs_client=True, needs_workspace=True)),
     "get_pipeline_logs": (get_pipeline_logs_tool, ToolConfig(needs_client=True, needs_workspace=True)),
     "search_pipeline": (search_pipeline_tool, ToolConfig(needs_client=True, needs_workspace=True)),
-    "list_indexes": (list_indexes_tool, ToolConfig(needs_client=True, needs_workspace=True)),
-    "get_index": (get_index_tool, ToolConfig(needs_client=True, needs_workspace=True)),
-    "create_index": (create_index_tool, ToolConfig(needs_client=True, needs_workspace=True)),
-    "update_index": (update_index_tool, ToolConfig(needs_client=True, needs_workspace=True)),
-    "deploy_index": (deploy_index_tool, ToolConfig(needs_client=True, needs_workspace=True)),
+    "list_indexes": (
+        list_indexes_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "get_index": (
+        get_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "create_index": (
+        create_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+    ),
+    "update_index": (
+        update_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.BOTH),
+    ),
+    "deploy_index": (
+        deploy_index_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
     "list_pipeline_templates": (list_pipeline_templates_tool, ToolConfig(needs_client=True, needs_workspace=True)),
     "get_pipeline_template": (get_pipeline_template_tool, ToolConfig(needs_client=True, needs_workspace=True)),
     "search_pipeline_templates": (
@@ -172,15 +196,15 @@ def create_enhanced_tool(
     """
     # Apply decorators first (if needed)
     decorated_func = base_func
-    if config.decorator_type != "none":
+    if config.memory_type != "none":
         store = STORE
         explorer = RichExplorer(store)
 
-        if config.decorator_type == "explorable":
+        if config.memory_type == "explorable":
             decorated_func = explorable(object_store=store, explorer=explorer)(decorated_func)
-        elif config.decorator_type == "referenceable":
+        elif config.memory_type == "referenceable":
             decorated_func = referenceable(object_store=store, explorer=explorer)(decorated_func)
-        elif config.decorator_type == "both":
+        elif config.memory_type == "both":
             decorated_func = explorable_and_referenceable(object_store=store, explorer=explorer)(decorated_func)
 
     # Handle client and workspace injection

src/deepset_mcp/tools/indexes.py

@@ -1,23 +1,36 @@
 from deepset_mcp.api.exceptions import BadRequestError, ResourceNotFoundError, UnexpectedAPIError
+from deepset_mcp.api.indexes.models import Index, IndexList
+from deepset_mcp.api.pipeline import PipelineValidationResult
 from deepset_mcp.api.protocols import AsyncClientProtocol
-from deepset_mcp.tools.formatting_utils import validation_result_to_llm_readable_string
-from deepset_mcp.tools.formatting_utils_index import index_list_to_llm_readable_string, index_to_llm_readable_string
 
 
-async def list_indexes(client: AsyncClientProtocol, workspace: str) -> str:
-    """Retrieves a list of all indexes available within the currently configured deepset workspace."""
-    response = await client.indexes(workspace=workspace).list()
-    return index_list_to_llm_readable_string(response)
+async def list_indexes(client: AsyncClientProtocol, workspace: str) -> IndexList | str:
+    """Use this to list available indexes on the deepset platform in your workspace.
+
+    :param client: Deepset API client to use for requesting indexes.
+    :param workspace: Workspace of which to list indexes.
+    """
+    try:
+        result = await client.indexes(workspace=workspace).list()
+    except ResourceNotFoundError as e:
+        return f"Error listing indexes. Error: {e.message} ({e.status_code})"
+
+    return result
 
 
-async def get_index(client: AsyncClientProtocol, workspace: str, index_name: str) -> str:
-    """Fetches detailed configuration information for a specific index, identified by its unique `index_name`."""
+async def get_index(client: AsyncClientProtocol, workspace: str, index_name: str) -> Index | str:
+    """Fetches detailed configuration information for a specific index, identified by its unique `index_name`.
+
+    :param client: Deepset API client to use for requesting the index.
+    :param workspace: Workspace of which to get the index from.
+    :param index_name: Unique name of the index to fetch.
+    """
     try:
         response = await client.indexes(workspace=workspace).get(index_name)
     except ResourceNotFoundError:
         return f"There is no index named '{index_name}'. Did you mean to create it?"
 
-    return index_to_llm_readable_string(response)
+    return response
 
 
 async def create_index(
@@ -26,10 +39,17 @@ async def create_index(
     index_name: str,
     yaml_configuration: str,
     description: str | None = None,
-) -> str:
-    """Creates a new index within the currently configured deepset workspace."""
+) -> dict[str, str | Index] | str:
+    """Creates a new index within your deepset platform workspace.
+
+    :param client: Deepset API client to use.
+    :param workspace: Workspace in which to create the index.
+    :param index_name: Unique name of the index to create.
+    :param yaml_configuration: YAML configuration to use for the index.
+    :param description: Description of the index to create.
+    """
     try:
-        await client.indexes(workspace=workspace).create(
+        result = await client.indexes(workspace=workspace).create(
             name=index_name, yaml_config=yaml_configuration, description=description
         )
     except ResourceNotFoundError:
@@ -39,7 +59,7 @@ async def create_index(
     except UnexpectedAPIError as e:
         return f"Failed to create index '{index_name}': {e}"
 
-    return f"Index '{index_name}' created successfully."
+    return {"message": f"Index '{index_name}' created successfully.", "index": result}
 
 
 async def update_index(
@@ -48,17 +68,23 @@ async def update_index(
     index_name: str,
     updated_index_name: str | None = None,
     yaml_configuration: str | None = None,
-) -> str:
-    """Updates an existing index in the specified workspace.
+) -> dict[str, str | Index] | str:
+    """Updates an existing index in your deepset platform workspace.
 
     This function can update either the name or the configuration of an existing index, or both.
     At least one of updated_index_name or yaml_configuration must be provided.
+
+    :param client: Deepset API client to use.
+    :param workspace: Workspace in which to update the index.
+    :param index_name: Unique name of the index to update.
+    :param updated_index_name: Updated name of the index.
+    :param yaml_configuration: YAML configuration to update the index with.
     """
     if not updated_index_name and not yaml_configuration:
         return "You must provide either a new name or a new configuration to update the index."
 
     try:
-        await client.indexes(workspace=workspace).update(
+        result = await client.indexes(workspace=workspace).update(
             index_name=index_name, updated_index_name=updated_index_name, yaml_config=yaml_configuration
         )
     except ResourceNotFoundError:
@@ -68,21 +94,21 @@ async def update_index(
     except UnexpectedAPIError as e:
         return f"Failed to update index '{index_name}': {e}"
 
-    return f"Index '{index_name}' updated successfully."
+    return {"message": f"Index '{index_name}' updated successfully.", "index": result}
 
 
-async def deploy_index(client: AsyncClientProtocol, workspace: str, index_name: str) -> str:
+async def deploy_index(client: AsyncClientProtocol, workspace: str, index_name: str) -> str | PipelineValidationResult:
     """Deploys an index to production.
 
     This function attempts to deploy the specified index in the given workspace.
-    If the deployment fails due to validation errors, it returns a readable string
+    If the deployment fails due to validation errors, it returns an object
     describing the validation errors.
 
     :param client: The async client for API communication.
     :param workspace: The workspace name.
     :param index_name: Name of the index to deploy.
 
-    :returns: A string indicating the deployment result.
+    :returns: A string indicating the deployment result or the validation results including errors.
     """
     try:
         deployment_result = await client.indexes(workspace=workspace).deploy(index_name=index_name)
@@ -94,6 +120,6 @@ async def deploy_index(client: AsyncClientProtocol, workspace: str, index_name:
         return f"Failed to deploy index '{index_name}': {e}"
 
     if not deployment_result.valid:
-        return validation_result_to_llm_readable_string(deployment_result)
+        return deployment_result
 
     return f"Index '{index_name}' deployed successfully."

test/unit/tools/test_indexes.py

@@ -19,6 +19,7 @@ def __init__(
         create_response: Index | None = None,
         update_response: Index | None = None,
         deploy_response: PipelineValidationResult | None = None,
+        list_exception: Exception | None = None,
         get_exception: Exception | None = None,
         create_exception: Exception | None = None,
         update_exception: Exception | None = None,
@@ -29,12 +30,15 @@ def __init__(
         self._create_response = create_response
         self._update_response = update_response
         self._deploy_response = deploy_response
+        self._list_exception = list_exception
         self._get_exception = get_exception
         self._create_exception = create_exception
         self._update_exception = update_exception
         self._deploy_exception = deploy_exception
 
     async def list(self, limit: int = 10, page_number: int = 1) -> IndexList:
+        if self._list_exception:
+            raise self._list_exception
         if self._list_response is not None:
             return self._list_response
         return IndexList(data=[], has_more=False, total=0)
@@ -118,17 +122,18 @@ def create_test_index(
 
 
 @pytest.mark.asyncio
-async def test_list_indexes_returns_formatted_string_when_no_indexes() -> None:
+async def test_list_indexes_without_indexes() -> None:
     resource = FakeIndexResource(list_response=IndexList(data=[], has_more=False, total=0))
     client = FakeClient(resource)
 
     result = await list_indexes(client=client, workspace="test")
 
-    assert result == "No indexes found."
+    assert isinstance(result, IndexList)
+    assert len(result.data) == 0
 
 
 @pytest.mark.asyncio
-async def test_list_indexes_returns_formatted_string_with_indexes() -> None:
+async def test_list_indexes_returns_indexes() -> None:
     index1 = create_test_index(name="index1", description="First index")
     index2 = create_test_index(name="index2", description="Second index")
 
@@ -137,25 +142,33 @@ async def test_list_indexes_returns_formatted_string_with_indexes() -> None:
 
     result = await list_indexes(client=client, workspace="test")
 
-    assert "index1" in result
-    assert "index2" in result
-    assert "First index" in result
-    assert "Second index" in result
-    assert "idx_123" in result
+    assert isinstance(result, IndexList)
+    assert len(result.data) == 2
+    assert result.data[0].name == index1.name
+    assert result.data[1].name == index2.name
 
 
 @pytest.mark.asyncio
-async def test_get_index_returns_formatted_string() -> None:
+async def test_list_indexes_returns_string_on_non_existant_workspace() -> None:
+    resource = FakeIndexResource(list_exception=ResourceNotFoundError(message="Resource not found."))
+    client = FakeClient(resource)
+
+    result = await list_indexes(client=client, workspace="test")
+
+    assert isinstance(result, str)
+    assert result == "Error listing indexes. Error: Resource not found. (404)"
+
+
+@pytest.mark.asyncio
+async def test_get_index_returns_index() -> None:
     index = create_test_index(name="my_index", description="My special index")
     resource = FakeIndexResource(get_response=index)
     client = FakeClient(resource)
 
     result = await get_index(client=client, workspace="test", index_name="my_index")
 
-    assert "my_index" in result
-    assert "config: value" in result
-    assert "idx_123" in result
-    assert "My special index" in result
+    assert isinstance(result, Index)
+    assert result.name == "my_index"
 
 
 @pytest.mark.asyncio
@@ -169,7 +182,7 @@ async def test_get_index_returns_error_message_when_index_not_found() -> None:
 
 
 @pytest.mark.asyncio
-async def test_create_index_returns_success_message() -> None:
+async def test_create_index_returns_success_message_and_index() -> None:
     created_index = create_test_index(name="new_index")
     resource = FakeIndexResource(create_response=created_index)
     client = FakeClient(resource)
@@ -182,7 +195,11 @@ async def test_create_index_returns_success_message() -> None:
         description="New index description",
     )
 
-    assert "Index 'new_index' created successfully." == result
+    assert isinstance(result, dict)
+    assert isinstance(result.get("message"), str)
+    index = result.get("index")
+    assert isinstance(index, Index)
+    assert index.name == "new_index"
 
 
 @pytest.mark.parametrize(
@@ -225,7 +242,12 @@ async def test_update_index_returns_success_message() -> None:
         yaml_configuration="new_config",
     )
 
-    assert "Index 'test_index' updated successfully." == result
+    assert isinstance(result, dict)
+    assert isinstance(result.get("message"), str)
+
+    index = result.get("index")
+    assert isinstance(index, Index)
+    assert index.name == "new_test_index"
 
 
 @pytest.mark.asyncio
@@ -355,11 +377,8 @@ async def test_deploy_index_returns_validation_errors() -> None:
 
     result = await deploy_index(client=client, workspace="test", index_name="test_index")
 
-    assert "The provided pipeline configuration is invalid" in result
-    assert "invalid_config" in result
-    assert "Index configuration is invalid" in result
-    assert "missing_dependency" in result
-    assert "Required dependency not found" in result
+    assert isinstance(result, PipelineValidationResult)
+    assert result.errors == validation_errors
 
 
 @pytest.mark.parametrize(