feat: use memory with pipeline template tools (#112)

Author: mathislucka

src/deepset_mcp/api/pipeline_template/models.py

@@ -31,3 +31,26 @@ class PipelineTemplate(BaseModel):
     yaml_config: str | None = Field(None, alias="query_yaml")
     tags: list[PipelineTemplateTag]
     pipeline_type: PipelineType
+
+
+class PipelineTemplateList(BaseModel):
+    """Response model for listing pipeline templates."""
+
+    data: list[PipelineTemplate]
+    has_more: bool
+    total: int
+
+
+class PipelineTemplateSearchResult(BaseModel):
+    """Model representing a search result for pipeline templates."""
+
+    template: PipelineTemplate
+    similarity_score: float
+
+
+class PipelineTemplateSearchResults(BaseModel):
+    """Response model for pipeline template search results."""
+
+    results: list[PipelineTemplateSearchResult]
+    query: str
+    total_found: int

src/deepset_mcp/api/pipeline_template/resource.py

@@ -1,7 +1,7 @@
 from typing import Any
 
 from deepset_mcp.api.exceptions import UnexpectedAPIError
-from deepset_mcp.api.pipeline_template.models import PipelineTemplate
+from deepset_mcp.api.pipeline_template.models import PipelineTemplate, PipelineTemplateList
 from deepset_mcp.api.protocols import AsyncClientProtocol
 from deepset_mcp.api.transport import raise_for_status
 
@@ -43,7 +43,7 @@ async def get_template(self, template_name: str) -> PipelineTemplate:
 
     async def list_templates(
         self, limit: int = 100, field: str = "created_at", order: str = "DESC", filter: str | None = None
-    ) -> list[PipelineTemplate]:
+    ) -> PipelineTemplateList:
         """List pipeline templates in the configured workspace.
 
         Parameters
@@ -59,8 +59,8 @@ async def list_templates(
 
         Returns
         -------
-        list[PipelineTemplate]
-            List of pipeline templates
+        PipelineTemplateList
+            List of pipeline templates with metadata
         """
         params = {"limit": limit, "page_number": 1, "field": field, "order": order}
 
@@ -80,4 +80,8 @@ async def list_templates(
 
         response_data: dict[str, Any] = response.json
 
-        return [PipelineTemplate.model_validate(template) for template in response_data["data"]]
+        return PipelineTemplateList(
+            data=[PipelineTemplate.model_validate(template) for template in response_data["data"]],
+            has_more=response_data.get("has_more", False),
+            total=response_data.get("total", len(response_data["data"])),
+        )

src/deepset_mcp/api/protocols.py

@@ -13,7 +13,7 @@
     PipelineLogList,
     PipelineValidationResult,
 )
-from deepset_mcp.api.pipeline_template.models import PipelineTemplate
+from deepset_mcp.api.pipeline_template.models import PipelineTemplate, PipelineTemplateList
 from deepset_mcp.api.secrets.models import Secret, SecretList
 from deepset_mcp.api.shared_models import DeepsetUser, NoContentResponse
 from deepset_mcp.api.transport import StreamingResponse, TransportResponse
@@ -237,7 +237,7 @@ async def get_template(self, template_name: str) -> PipelineTemplate:
 
     async def list_templates(
         self, limit: int = 100, field: str = "created_at", order: str = "DESC", filter: str | None = None
-    ) -> list[PipelineTemplate]:
+    ) -> PipelineTemplateList:
         """List pipeline templates in the configured workspace."""
         ...
 

src/deepset_mcp/tool_factory.py

@@ -75,7 +75,7 @@ async def search_component_definitions_wrapper(client: Any, query: str, top_k: i
 
 
 # Special wrapper for search_pipeline_templates that needs the model
-async def search_pipeline_templates_wrapper(client: Any, workspace: str, query: str, top_k: int = 10) -> str:
+async def search_pipeline_templates_wrapper(client: Any, workspace: str, query: str, top_k: int = 10) -> Any:
     """Searches for pipeline templates based on name or description using semantic similarity.
 
     Args:
@@ -85,7 +85,7 @@ async def search_pipeline_templates_wrapper(client: Any, workspace: str, query:
         top_k: Maximum number of results to return (default: 10)
 
     Returns:
-        A formatted string containing the matched pipeline template definitions
+        Search results with similarity scores or error message
     """
     model = get_initialized_model()
     return await search_pipeline_templates_tool(client, query, model, workspace, top_k)
@@ -156,11 +156,17 @@ def get_workspace_from_env() -> str:
         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)),
+    "list_pipeline_templates": (
+        list_pipeline_templates_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
+    "get_pipeline_template": (
+        get_pipeline_template_tool,
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
+    ),
     "search_pipeline_templates": (
         search_pipeline_templates_wrapper,
-        ToolConfig(needs_client=True, needs_workspace=True),
+        ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
     ),
     "list_custom_component_installations": (
         list_custom_component_installations_tool,

src/deepset_mcp/tools/pipeline_template.py

@@ -1,8 +1,13 @@
 import numpy as np
 
 from deepset_mcp.api.exceptions import ResourceNotFoundError, UnexpectedAPIError
+from deepset_mcp.api.pipeline_template.models import (
+    PipelineTemplate,
+    PipelineTemplateList,
+    PipelineTemplateSearchResult,
+    PipelineTemplateSearchResults,
+)
 from deepset_mcp.api.protocols import AsyncClientProtocol
-from deepset_mcp.tools.formatting_utils import pipeline_template_to_llm_readable_string
 from deepset_mcp.tools.model_protocol import ModelProtocol
 
 
@@ -13,7 +18,7 @@ async def list_pipeline_templates(
     field: str = "created_at",
     order: str = "DESC",
     filter: str | None = None,
-) -> str:
+) -> PipelineTemplateList | str:
     """Retrieves a list of all available pipeline templates.
 
     :param client: The async client for API requests.
@@ -23,25 +28,31 @@ async def list_pipeline_templates(
     :param order: Sort order, either "ASC" or "DESC" (default: "DESC").
     :param filter: OData filter expression to filter templates by criteria.
 
-    :returns: Formatted string with template information.
+    :returns: List of pipeline templates or error message.
     """
     try:
-        response = await client.pipeline_templates(workspace=workspace).list_templates(
+        return await client.pipeline_templates(workspace=workspace).list_templates(
             limit=limit, field=field, order=order, filter=filter
         )
-        formatted_templates = [pipeline_template_to_llm_readable_string(t) for t in response]
-        return "\n\n".join(formatted_templates)
     except ResourceNotFoundError:
         return f"There is no workspace named '{workspace}'. Did you mean to configure it?"
     except UnexpectedAPIError as e:
         return f"Failed to list pipeline templates: {e}"
 
 
-async def get_pipeline_template(client: AsyncClientProtocol, workspace: str, template_name: str) -> str:
-    """Fetches detailed information for a specific pipeline template, identified by its `template_name`."""
+async def get_pipeline_template(
+    client: AsyncClientProtocol, workspace: str, template_name: str
+) -> PipelineTemplate | str:
+    """Fetches detailed information for a specific pipeline template, identified by its `template_name`.
+
+    :param client: The async client for API requests.
+    :param workspace: The workspace to fetch template from.
+    :param template_name: The name of the template to fetch.
+
+    :returns: Pipeline template details or error message.
+    """
     try:
-        response = await client.pipeline_templates(workspace=workspace).get_template(template_name)
-        return pipeline_template_to_llm_readable_string(response, include_yaml=True)
+        return await client.pipeline_templates(workspace=workspace).get_template(template_name)
     except ResourceNotFoundError:
         return f"There is no pipeline template named '{template_name}' in workspace '{workspace}'."
     except UnexpectedAPIError as e:
@@ -50,18 +61,16 @@ async def get_pipeline_template(client: AsyncClientProtocol, workspace: str, tem
 
 async def search_pipeline_templates(
     client: AsyncClientProtocol, query: str, model: ModelProtocol, workspace: str, top_k: int = 10
-) -> str:
+) -> PipelineTemplateSearchResults | str:
     """Searches for pipeline templates based on name or description using semantic similarity.
 
-    Args:
-        client: The API client to use
-        query: The search query
-        model: The model to use for computing embeddings
-        workspace: The workspace to search templates from
-        top_k: Maximum number of results to return (default: 5)
+    :param client: The API client to use.
+    :param query: The search query.
+    :param model: The model to use for computing embeddings.
+    :param workspace: The workspace to search templates from.
+    :param top_k: Maximum number of results to return (default: 10).
 
-    Returns:
-        A formatted string containing the matched pipeline template definitions
+    :returns: Search results with similarity scores or error message.
     """
     try:
         response = await client.pipeline_templates(workspace=workspace).list_templates(
@@ -70,12 +79,12 @@ async def search_pipeline_templates(
     except UnexpectedAPIError as e:
         return f"Failed to retrieve pipeline templates: {e}"
 
-    if not response:
-        return "No pipeline templates found"
+    if not response.data:
+        return PipelineTemplateSearchResults(results=[], query=query, total_found=0)
 
     # Extract text for embedding from all templates
     template_texts: list[tuple[str, str]] = [
-        (template.template_name, f"{template.template_name} {template.description}") for template in response
+        (template.template_name, f"{template.template_name} {template.description}") for template in response.data
     ]
     template_names: list[str] = [t[0] for t in template_texts]
 
@@ -96,12 +105,11 @@ async def search_pipeline_templates(
     template_similarities.sort(key=lambda x: x[1], reverse=True)
 
     top_templates = template_similarities[:top_k]
-    results = []
+    search_results = []
     for template_name, sim in top_templates:
         # Find the template object by name
-        template = next((t for t in response if t.template_name == template_name), None)
+        template = next((t for t in response.data if t.template_name == template_name), None)
         if template:
-            template_str = pipeline_template_to_llm_readable_string(template)
-            results.append(f"Similarity Score: {sim:.3f}\n{template_str}\n{'-' * 80}\n")
+            search_results.append(PipelineTemplateSearchResult(template=template, similarity_score=float(sim)))
 
-    return "\n".join(results)
+    return PipelineTemplateSearchResults(results=search_results, query=query, total_found=len(search_results))

test/integration/test_integration_pipeline_template_resource.py

@@ -2,7 +2,7 @@
 
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.api.exceptions import ResourceNotFoundError
-from deepset_mcp.api.pipeline_template.models import PipelineTemplate
+from deepset_mcp.api.pipeline_template.models import PipelineTemplate, PipelineTemplateList
 from deepset_mcp.api.pipeline_template.resource import PipelineTemplateResource
 
 pytestmark = pytest.mark.integration
@@ -26,14 +26,14 @@ async def test_get_template(
     First lists all templates, then gets the first one by name.
     """
     # Get all templates to find an existing one
-    templates = await template_resource.list_templates()
+    templates_list = await template_resource.list_templates()
 
     # Skip if no templates are available
-    if not templates:
+    if not templates_list.data:
         pytest.skip("No templates available in the test environment")
 
     # Get the first template's name
-    template_name = templates[0].template_name
+    template_name = templates_list.data[0].template_name
 
     # Now get that specific template
     template = await template_resource.get_template(template_name=template_name)
@@ -64,17 +64,18 @@ async def test_list_templates(
 ) -> None:
     """Test listing templates."""
     # Test listing templates with default limit
-    templates = await template_resource.list_templates()
+    templates_list = await template_resource.list_templates()
 
-    # Verify that the templates are returned as a list
-    assert isinstance(templates, list)
+    # Verify that the templates are returned as a PipelineTemplateList
+    assert isinstance(templates_list, PipelineTemplateList)
+    assert isinstance(templates_list.data, list)
 
     # Skip further checks if no templates are available
-    if not templates:
+    if not templates_list.data:
         pytest.skip("No templates available in the test environment")
 
     # Verify the first template has the expected structure
-    template = templates[0]
+    template = templates_list.data[0]
     assert isinstance(template, PipelineTemplate)
     assert template.template_name is not None
     assert template.author is not None
@@ -89,10 +90,11 @@ async def test_list_templates_with_limit(
     """Test listing templates with a specific limit."""
     # Test with a small limit
     limit = 1
-    templates = await template_resource.list_templates(limit=limit)
+    templates_list = await template_resource.list_templates(limit=limit)
 
     # Verify that the number of templates is not more than the limit
-    assert len(templates) <= limit
+    assert isinstance(templates_list, PipelineTemplateList)
+    assert len(templates_list.data) <= limit
 
 
 @pytest.mark.asyncio
@@ -101,24 +103,26 @@ async def test_list_templates_with_filter(
 ) -> None:
     """Test listing templates with a pipeline type filter."""
     # Test filtering by QUERY pipeline type
-    query_templates = await template_resource.list_templates(filter="pipeline_type eq 'QUERY'")
+    query_templates_list = await template_resource.list_templates(filter="pipeline_type eq 'QUERY'")
 
     # Verify that all returned templates are QUERY type
-    assert isinstance(query_templates, list)
+    assert isinstance(query_templates_list, PipelineTemplateList)
+    assert isinstance(query_templates_list.data, list)
 
     # If templates are available, verify they are all QUERY type
-    for template in query_templates:
+    for template in query_templates_list.data:
         assert isinstance(template, PipelineTemplate)
         assert template.pipeline_type == "query"
 
     # Test filtering by INDEXING pipeline type
-    indexing_templates = await template_resource.list_templates(filter="pipeline_type eq 'INDEXING'")
+    indexing_templates_list = await template_resource.list_templates(filter="pipeline_type eq 'INDEXING'")
 
     # Verify that all returned templates are INDEXING type
-    assert isinstance(indexing_templates, list)
+    assert isinstance(indexing_templates_list, PipelineTemplateList)
+    assert isinstance(indexing_templates_list.data, list)
 
     # If templates are available, verify they are all INDEXING type
-    for template in indexing_templates:
+    for template in indexing_templates_list.data:
         assert isinstance(template, PipelineTemplate)
         assert template.pipeline_type == "indexing"
 
@@ -129,12 +133,13 @@ async def test_list_templates_with_custom_sorting(
 ) -> None:
     """Test listing templates with custom sorting."""
     # Test sorting by name in ascending order
-    templates = await template_resource.list_templates(field="name", order="ASC", limit=5)
+    templates_list = await template_resource.list_templates(field="name", order="ASC", limit=5)
 
-    # Verify that the templates are returned as a list
-    assert isinstance(templates, list)
+    # Verify that the templates are returned as a PipelineTemplateList
+    assert isinstance(templates_list, PipelineTemplateList)
+    assert isinstance(templates_list.data, list)
 
     # If we have multiple templates, verify they are sorted correctly
-    if len(templates) > 1:
-        for i in range(len(templates) - 1):
-            assert templates[i].display_name <= templates[i + 1].display_name
+    if len(templates_list.data) > 1:
+        for i in range(len(templates_list.data) - 1):
+            assert templates_list.data[i].display_name <= templates_list.data[i + 1].display_name