feat: Add search_pipeline_templates tool with semantic similarity search (#59)

* feat: add import for pipeline template formatting utility

* feat: add search_pipeline_templates function with semantic similarity search

* feat: add import for search_pipeline_templates function

* feat: add import for pipeline template models for testing

* feat: add fake pipeline templates resource for testing

* feat: extend fake client to support pipeline templates resource

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* fix: use keyword arguments in fake client initialization

* feat: add uuid import for test data generation

* feat: add more embedding patterns for pipeline template testing

* feat: add unit tests for search_pipeline_templates function

* feat: import search_pipeline_templates tool in main module

* feat: add search_pipeline_templates tool to main MCP server

* feat: add test runner to verify our implementation

* feat: create simpler test runner for verification

* fix: handle import issues in test runner

* feat: create simple import checker

* chore: remove test runner as it's not needed for this PR

* chore: remove import checker as it's not needed for this PR

* refactor: remove search_pipeline_templates from haystack_service.py

* feat: add imports for search_pipeline_templates function

* feat: add search_pipeline_templates function to pipeline_template.py

* fix: remove search_pipeline_templates import from haystack_service

* feat: import search_pipeline_templates from pipeline_template module

* test: remove search_pipeline_templates import from haystack_service tests

* test: remove unused PipelineTemplate imports from haystack_service tests

* test: remove FakePipelineTemplatesResource from haystack_service tests

* test: simplify FakeClient to remove pipeline_templates functionality

* test: remove search_pipeline_templates tests from haystack_service tests

* test: add imports for search_pipeline_templates tests

* test: add FakeModel class for search_pipeline_templates tests

* test: add search_pipeline_templates unit tests

* test: fix PipelineTemplate instantiation with correct field names

* fix: lint

Author: mathislucka

src/deepset_mcp/main.py

@@ -27,6 +27,7 @@
 from deepset_mcp.tools.pipeline_template import (
     get_pipeline_template as get_pipeline_template_tool,
     list_pipeline_templates as list_pipeline_templates_tool,
+    search_pipeline_templates as search_pipeline_templates_tool,
 )
 
 INITIALIZED_MODEL = StaticModel.from_pretrained("minishlab/potion-base-2M")
@@ -202,6 +203,24 @@ async def search_component_definitions(query: str) -> str:
     return response
 
 
+@mcp.tool()
+async def search_pipeline_templates(query: str) -> str:
+    """Use this to search for pipeline templates in deepset.
+
+    You can use full natural language queries to find templates.
+    You can also use simple keywords.
+    Use this if you want to find pipeline templates for specific use cases,
+    but you are not sure what the exact name of the template is.
+    """
+    workspace = get_workspace()
+    async with AsyncDeepsetClient() as client:
+        response = await search_pipeline_templates_tool(
+            client=client, query=query, model=INITIALIZED_MODEL, workspace=workspace
+        )
+
+    return response
+
+
 @mcp.tool()
 async def list_indexes() -> str:
     """Retrieves a list of all indexes available in the deepset workspace.

src/deepset_mcp/tools/pipeline_template.py

@@ -1,6 +1,9 @@
+import numpy as np
+
 from deepset_mcp.api.exceptions import ResourceNotFoundError, UnexpectedAPIError
 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
 
 
 async def list_pipeline_templates(
@@ -43,3 +46,60 @@ async def get_pipeline_template(client: AsyncClientProtocol, workspace: str, tem
         return f"There is no pipeline template named '{template_name}' in workspace '{workspace}'."
     except UnexpectedAPIError as e:
         return f"Failed to fetch pipeline template '{template_name}': {e}"
+
+
+async def search_pipeline_templates(
+    client: AsyncClientProtocol, query: str, model: ModelProtocol, workspace: str, top_k: int = 5
+) -> 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)
+
+    Returns:
+        A formatted string containing the matched pipeline template definitions
+    """
+    try:
+        response = await client.pipeline_templates(workspace=workspace).list_templates()
+    except UnexpectedAPIError as e:
+        return f"Failed to retrieve pipeline templates: {e}"
+
+    if not response:
+        return "No pipeline templates found"
+
+    # 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_names: list[str] = [t[0] for t in template_texts]
+
+    # Compute embeddings
+    query_embedding = model.encode(query)
+    template_embeddings = model.encode([text for _, text in template_texts])
+
+    query_embedding_reshaped = query_embedding.reshape(1, -1)
+
+    # Calculate dot product between target and all templates
+    # This gives us a similarity score for each template
+    similarities = np.dot(template_embeddings, query_embedding_reshaped.T).flatten()
+
+    # Create (template_name, similarity) pairs
+    template_similarities = list(zip(template_names, similarities, strict=False))
+
+    # Sort by similarity score in descending order
+    template_similarities.sort(key=lambda x: x[1], reverse=True)
+
+    top_templates = template_similarities[:top_k]
+    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)
+        if template:
+            template_str = pipeline_template_to_llm_readable_string(template)
+            results.append(f"Similarity Score: {sim:.3f}\n{template_str}\n{'-' * 80}\n")
+
+    return "\n".join(results)

test/unit/tools/test_haystack_service.py

@@ -26,6 +26,10 @@ def encode(self, sentences: list[str] | str) -> np.ndarray[Any, Any]:
                 embeddings[i] = [0, 0, 0.9]
             elif "reader" in sentence.lower():
                 embeddings[i] = [0, 1, 0]
+            elif "rag" in sentence.lower() or "retrieval" in sentence.lower():
+                embeddings[i] = [1, 0, 0]
+            elif "chat" in sentence.lower() or "conversation" in sentence.lower():
+                embeddings[i] = [0.8, 0.2, 0]
             else:
                 embeddings[i] = [0, 0, 1]
         return embeddings
@@ -58,11 +62,16 @@ async def get_component_input_output(self, component_name: str) -> dict[str, Any
 
 
 class FakeClient(BaseFakeClient):
-    def __init__(self, resource: FakeHaystackServiceResource):
+    def __init__(
+        self,
+        resource: FakeHaystackServiceResource | None = None,
+    ):
         self._resource = resource
         super().__init__()
 
     def haystack_service(self) -> FakeHaystackServiceResource:
+        if self._resource is None:
+            raise ValueError("Haystack service resource not configured")
         return self._resource
 
 
@@ -139,7 +148,7 @@ async def test_get_component_definition_success() -> None:
     resource = FakeHaystackServiceResource(
         get_component_schemas_response=schema_response, get_component_io_response=io_response
     )
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     result = await get_component_definition(client, component_type)
 
     # Check that all required information is present
@@ -169,7 +178,7 @@ async def test_get_component_definition_success() -> None:
 async def test_get_component_definition_not_found() -> None:
     response: dict[str, Any] = {"component_schema": {"definitions": {"Components": {}}}}
     resource = FakeHaystackServiceResource(get_component_schemas_response=response)
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     result = await get_component_definition(client, "nonexistent.component")
     assert "Component not found" in result
 
@@ -213,7 +222,7 @@ async def test_search_component_definition_success() -> None:
     resource = FakeHaystackServiceResource(
         get_component_schemas_response=schema_response, get_component_io_response=io_response
     )
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     model = FakeModel()
 
     # Search for converters
@@ -232,7 +241,7 @@ async def test_search_component_definition_success() -> None:
 @pytest.mark.asyncio
 async def test_get_component_definition_api_error() -> None:
     resource = FakeHaystackServiceResource(exception=UnexpectedAPIError(status_code=500, message="API Error"))
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     result = await get_component_definition(client, "some.component")
     assert "Failed to retrieve component definition" in result
     assert "API Error" in result
@@ -242,7 +251,7 @@ async def test_get_component_definition_api_error() -> None:
 async def test_search_component_definition_no_components() -> None:
     schema_response: dict[str, Any] = {"component_schema": {"definitions": {"Components": {}}}}
     resource = FakeHaystackServiceResource(get_component_schemas_response=schema_response)
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     model = FakeModel()
 
     result = await search_component_definition(client, "test query", model)
@@ -252,7 +261,7 @@ async def test_search_component_definition_no_components() -> None:
 @pytest.mark.asyncio
 async def test_search_component_definition_api_error() -> None:
     resource = FakeHaystackServiceResource(exception=UnexpectedAPIError(status_code=500, message="API Error"))
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     model = FakeModel()
 
     result = await search_component_definition(client, "test query", model)
@@ -263,7 +272,7 @@ async def test_search_component_definition_api_error() -> None:
 async def test_list_component_families_no_families() -> None:
     response: dict[str, Any] = {"component_schema": {"definitions": {"Components": {}}}}
     resource = FakeHaystackServiceResource(get_component_schemas_response=response)
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     result = await list_component_families(client)
     assert "No component families found" in result
 
@@ -287,7 +296,7 @@ async def test_list_component_families_success() -> None:
         }
     }
     resource = FakeHaystackServiceResource(get_component_schemas_response=response)
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     result = await list_component_families(client)
 
     assert "Available Haystack component families" in result
@@ -302,7 +311,7 @@ async def test_list_component_families_success() -> None:
 @pytest.mark.asyncio
 async def test_list_component_families_api_error() -> None:
     resource = FakeHaystackServiceResource(exception=UnexpectedAPIError(status_code=500, message="API Error"))
-    client = FakeClient(resource)
+    client = FakeClient(resource=resource)
     result = await list_component_families(client)
     assert "Failed to retrieve component families" in result
     assert "API Error" in result

test/unit/tools/test_pipeline_template.py

@@ -1,14 +1,38 @@
 from typing import Any
-from uuid import UUID
+from uuid import UUID, uuid4
 
+import numpy as np
 import pytest
 
 from deepset_mcp.api.exceptions import ResourceNotFoundError, UnexpectedAPIError
 from deepset_mcp.api.pipeline_template.models import PipelineTemplate, PipelineTemplateTag, PipelineType
-from deepset_mcp.tools.pipeline_template import get_pipeline_template, list_pipeline_templates
+from deepset_mcp.tools.model_protocol import ModelProtocol
+from deepset_mcp.tools.pipeline_template import (
+    get_pipeline_template,
+    list_pipeline_templates,
+    search_pipeline_templates,
+)
 from test.unit.conftest import BaseFakeClient
 
 
+class FakeModel(ModelProtocol):
+    def encode(self, sentences: list[str] | str) -> np.ndarray[Any, Any]:
+        # Convert input to list if it's a single string
+        if isinstance(sentences, str):
+            sentences = [sentences]
+
+        # Create fake embeddings with consistent similarities
+        embeddings = np.zeros((len(sentences), 3))
+        for i, sentence in enumerate(sentences):
+            if "rag" in sentence.lower() or "retrieval" in sentence.lower():
+                embeddings[i] = [1, 0, 0]
+            elif "chat" in sentence.lower() or "conversation" in sentence.lower():
+                embeddings[i] = [0.8, 0.2, 0]
+            else:
+                embeddings[i] = [0, 0, 1]
+        return embeddings
+
+
 class FakePipelineTemplateResource:
     def __init__(
         self,
@@ -231,3 +255,70 @@ async def test_list_pipeline_templates_with_filter_and_sorting() -> None:
     assert resource.last_list_call_params["field"] == "name"
     assert resource.last_list_call_params["order"] == "ASC"
     assert resource.last_list_call_params["filter"] == filter_value
+
+
+@pytest.mark.asyncio
+async def test_search_pipeline_templates_success() -> None:
+    # Create sample pipeline templates
+    templates = [
+        PipelineTemplate(
+            author="Deepset",
+            best_for=["Document Q&A"],
+            description="A retrieval-augmented generation template for answering questions",
+            pipeline_name="rag-pipeline",
+            name="RAG Pipeline",
+            pipeline_template_id=uuid4(),
+            potential_applications=["FAQ systems", "Document search"],
+            query_yaml="components:\n  retriever: ...\n  generator: ...",
+            tags=[],
+            pipeline_type=PipelineType.QUERY,
+        ),
+        PipelineTemplate(
+            author="Deepset",
+            best_for=["Conversational AI"],
+            description="A chat-based conversational pipeline for interactive responses",
+            pipeline_name="chat-pipeline",
+            name="Chat Pipeline",
+            pipeline_template_id=uuid4(),
+            potential_applications=["Chatbots", "Virtual assistants"],
+            query_yaml="components:\n  chat_generator: ...\n  memory: ...",
+            tags=[],
+            pipeline_type=PipelineType.QUERY,
+        ),
+    ]
+
+    resource = FakePipelineTemplateResource(list_response=templates)
+    client = FakeClient(resource)
+    model = FakeModel()
+
+    # Search for RAG templates
+    result = await search_pipeline_templates(client, "retrieval augmented generation", model, "test_workspace")
+    assert "rag-pipeline" in result
+    assert "Similarity Score:" in result
+    assert "retrieval-augmented generation" in result
+
+    # Search for chat templates
+    result = await search_pipeline_templates(client, "conversational chat interface", model, "test_workspace")
+    assert "chat-pipeline" in result
+    assert "Similarity Score:" in result
+    assert "chat-based conversational" in result
+
+
+@pytest.mark.asyncio
+async def test_search_pipeline_templates_no_templates() -> None:
+    resource = FakePipelineTemplateResource(list_response=[])
+    client = FakeClient(resource)
+    model = FakeModel()
+
+    result = await search_pipeline_templates(client, "test query", model, "test_workspace")
+    assert "No pipeline templates found" in result
+
+
+@pytest.mark.asyncio
+async def test_search_pipeline_templates_api_error() -> None:
+    resource = FakePipelineTemplateResource(list_exception=UnexpectedAPIError(status_code=500, message="API Error"))
+    client = FakeClient(resource)
+    model = FakeModel()
+
+    result = await search_pipeline_templates(client, "test query", model, "test_workspace")
+    assert "Failed to retrieve pipeline templates" in result