feat: add semantic search for component definitions (#33)

* feat: add model protocol for embeddings

* feat: add component helper functions

* feat: update imports for search functionality

* refactor: move component info extraction to helper function

* refactor: move io info formatting to helper function

* feat: add search_component_definition function

* test: add tests for component helper functions

* test: update imports for search functionality tests

* test: add fake model implementation

* test: add io response parameter to fake service

* test: use separate response for io information

* test: update component definition test to use updated fake service

* test: add success case for search component functionality

* test: add error cases for search component functionality

* feat: component search mcp

* fix: linting and types

Author: mathislucka

pyproject.toml

@@ -12,6 +12,8 @@ dependencies = [
     "httpx",
     "pydantic>=2.0.0",
     "pyyaml",
+    "numpy",
+    "model2vec",
 ]
 
 [project.scripts]

src/deepset_mcp/main.py

@@ -1,11 +1,13 @@
 import os
 
 from mcp.server.fastmcp import FastMCP
+from model2vec import StaticModel
 
 from deepset_mcp.api.client import AsyncDeepsetClient
 from deepset_mcp.tools.haystack_service import (
     get_component_definition as get_component_definition_tool,
     list_component_families as list_component_families_tool,
+    search_component_definition as search_component_definition_tool,
 )
 from deepset_mcp.tools.pipeline import (
     create_pipeline as create_pipeline_tool,
@@ -15,6 +17,8 @@
     validate_pipeline as validate_pipeline_tool,
 )
 
+INITIALIZED_MODEL = StaticModel.from_pretrained("minishlab/potion-base-2M")
+
 # Initialize MCP Server
 mcp = FastMCP("Deepset Cloud MCP")
 
@@ -134,6 +138,21 @@ async def validate_pipeline(yaml_configuration: str) -> str:
     return response
 
 
+@mcp.tool()
+async def search_component_definitions(query: str) -> str:
+    """Use this to search for components in deepset.
+
+    You can use full natural language queries to find components.
+    You can also use simple keywords.
+    Use this if you want to find the definition for a component,
+    but you are not sure what the exact name of the component is.
+    """
+    async with AsyncDeepsetClient() as client:
+        response = await search_component_definition_tool(client=client, query=query, model=INITIALIZED_MODEL)
+
+    return response
+
+
 #
 #
 # @mcp.tool()

src/deepset_mcp/tools/component_helper.py

@@ -0,0 +1,123 @@
+from typing import Any
+
+
+def extract_component_info(components: dict[str, Any], component_def: dict[str, Any]) -> str:
+    """Extracts and formats component information from its definition.
+
+    Args:
+        components: The components dictionary from the schema
+        component_def: The specific component definition
+
+    Returns:
+        A formatted string containing the component information
+    """
+    component_type_info = component_def["properties"]["type"]
+    init_params = component_def["properties"].get("init_parameters", {}).get("properties", {})
+    component_type = component_type_info["const"]
+
+    # Format the basic component information
+    parts = [
+        f"Component: {component_type}",
+        f"Name: {component_def.get('title', 'Unknown')}",
+        f"Family: {component_type_info.get('family', 'Unknown')}",
+        f"Family Description: {component_type_info.get('family_description', 'No description available.')}",
+        f"\nDescription:\n{component_def.get('description', 'No description available.')}\n",
+        "\nInitialization Parameters:",
+    ]
+
+    if not init_params:
+        parts.append("  No initialization parameters")
+    else:
+        for param_name, param_info in init_params.items():
+            param_type = param_info.get("_annotation", param_info.get("type", "Unknown"))
+            param_desc = param_info.get("description", "No description available.")
+            default = f" (default: {param_info['default']})" if "default" in param_info else ""
+            parts.append(f"  {param_name}: {param_type}{default}\n    {param_desc}")
+
+    return "\n".join(parts)
+
+
+def format_io_info(io_info: dict[str, Any]) -> str:
+    """Formats the input/output information for a component.
+
+    Args:
+        io_info: The input/output information dictionary
+
+    Returns:
+        A formatted string containing the IO information
+    """
+    parts = []
+
+    # Add Input Schema
+    parts.append("\nInput Schema:")
+    if "input" in io_info:
+        input_props = io_info["input"].get("properties", {})
+        if not input_props:
+            parts.append("  No input parameters")
+        else:
+            required = io_info["input"].get("required", [])
+            for param_name, param_info in input_props.items():
+                req_marker = " (required)" if param_name in required else ""
+                param_type = param_info.get("_annotation", param_info.get("type", "Unknown"))
+                param_desc = param_info.get("description", "No description available.")
+                default = f" (default: {param_info['default']})" if "default" in param_info else ""
+                parts.append(f"  {param_name}: {param_type}{req_marker}{default}\n    {param_desc}")
+    else:
+        parts.append("  Input schema not available")
+
+    # Add Output Schema
+    parts.append("\nOutput Schema:")
+    if "output" in io_info and isinstance(io_info["output"], dict):
+        output_info = io_info["output"]
+        if "properties" in output_info:
+            output_props = output_info.get("properties", {})
+            if not output_props:
+                parts.append("  No output parameters")
+            else:
+                required = output_info.get("required", [])
+                for param_name, param_info in output_props.items():
+                    req_marker = " (required)" if param_name in required else ""
+                    param_type = param_info.get("_annotation", param_info.get("type", "Unknown"))
+                    param_desc = param_info.get("description", "No description available.")
+                    default = f" (default: {param_info['default']})" if "default" in param_info else ""
+                    parts.append(f"  {param_name}: {param_type}{req_marker}{default}\n    {param_desc}")
+
+                # Include any definitions if they exist
+                if "definitions" in output_info:
+                    parts.append("\n  Definitions:")
+                    for def_name, def_info in output_info["definitions"].items():
+                        parts.append(f"\n    {def_name}:")
+                        if "properties" in def_info:
+                            def_required = def_info.get("required", [])
+                            for prop_name, prop_info in def_info["properties"].items():
+                                req_marker = " (required)" if prop_name in def_required else ""
+                                prop_type = prop_info.get("_annotation", prop_info.get("type", "Unknown"))
+                                prop_desc = prop_info.get("description", "No description available.")
+                                default = f" (default: {prop_info['default']})" if "default" in prop_info else ""
+                                parts.append(
+                                    f"      {prop_name}: {prop_type}{req_marker}{default}\n        {prop_desc}"
+                                )
+        else:
+            # Simple output schema
+            desc = output_info.get("description", "No description available.")
+            output_type = output_info.get("type", "Unknown")
+            parts.append(f"  Type: {output_type}\n  {desc}")
+    else:
+        parts.append("  Output schema not available")
+
+    return "\n".join(parts)
+
+
+def extract_component_texts(component_def: dict[str, Any]) -> tuple[str, str]:
+    """Extracts the component name and description for embedding.
+
+    Args:
+        component_def: The component definition
+
+    Returns:
+        A tuple containing the component name and description
+    """
+    component_type = component_def["properties"]["type"]["const"]
+    name = component_def.get("title", "")
+    description = component_def.get("description", "")
+    return component_type, f"{name} {description}"

src/deepset_mcp/tools/haystack_service.py

@@ -1,5 +1,13 @@
+import numpy as np
+
 from deepset_mcp.api.exceptions import UnexpectedAPIError
 from deepset_mcp.api.protocols import AsyncClientProtocol
+from deepset_mcp.tools.component_helper import (
+    extract_component_info,
+    extract_component_texts,
+    format_io_info,
+)
+from deepset_mcp.tools.model_protocol import ModelProtocol
 
 
 async def get_component_definition(client: AsyncClientProtocol, component_type: str) -> str:
@@ -32,97 +40,76 @@ async def get_component_definition(client: AsyncClientProtocol, component_type:
     if not component_def:
         return f"Component not found: {component_type}"
 
-    # Extract relevant information
-    component_type_info = component_def["properties"]["type"]
-    init_params = component_def["properties"].get("init_parameters", {}).get("properties", {})
-
-    # Format the basic component information
-    parts = [
-        f"Component: {component_type}",
-        f"Name: {component_def.get('title', 'Unknown')}",
-        f"Family: {component_type_info.get('family', 'Unknown')}",
-        f"Family Description: {component_type_info.get('family_description', 'No description available.')}",
-        f"\nDescription:\n{component_def.get('description', 'No description available.')}\n",
-        "\nInitialization Parameters:",
-    ]
-
-    if not init_params:
-        parts.append("  No initialization parameters")
-    else:
-        for param_name, param_info in init_params.items():
-            param_type = param_info.get("_annotation", param_info.get("type", "Unknown"))
-            param_desc = param_info.get("description", "No description available.")
-            default = f" (default: {param_info['default']})" if "default" in param_info else ""
-            parts.append(f"  {param_name}: {param_type}{default}\n    {param_desc}")
+    # Get component information
+    parts = [extract_component_info(components, component_def)]
 
     # Fetch and add input/output information
     try:
         # Extract component name from the full path
         component_name = component_type.split(".")[-1]
         io_info = await haystack_service.get_component_input_output(component_name)
-
-        # Add Input Schema
-        parts.append("\nInput Schema:")
-        if "input" in io_info:
-            input_props = io_info["input"].get("properties", {})
-            if not input_props:
-                parts.append("  No input parameters")
-            else:
-                required = io_info["input"].get("required", [])
-                for param_name, param_info in input_props.items():
-                    req_marker = " (required)" if param_name in required else ""
-                    param_type = param_info.get("_annotation", param_info.get("type", "Unknown"))
-                    param_desc = param_info.get("description", "No description available.")
-                    default = f" (default: {param_info['default']})" if "default" in param_info else ""
-                    parts.append(f"  {param_name}: {param_type}{req_marker}{default}\n    {param_desc}")
-        else:
-            parts.append("  Input schema not available")
-
-        # Add Output Schema
-        parts.append("\nOutput Schema:")
-        if "output" in io_info and isinstance(io_info["output"], dict):
-            output_info = io_info["output"]
-            if "properties" in output_info:
-                output_props = output_info.get("properties", {})
-                if not output_props:
-                    parts.append("  No output parameters")
-                else:
-                    required = output_info.get("required", [])
-                    for param_name, param_info in output_props.items():
-                        req_marker = " (required)" if param_name in required else ""
-                        param_type = param_info.get("_annotation", param_info.get("type", "Unknown"))
-                        param_desc = param_info.get("description", "No description available.")
-                        default = f" (default: {param_info['default']})" if "default" in param_info else ""
-                        parts.append(f"  {param_name}: {param_type}{req_marker}{default}\n    {param_desc}")
-
-                    # Include any definitions if they exist
-                    if "definitions" in output_info:
-                        parts.append("\n  Definitions:")
-                        for def_name, def_info in output_info["definitions"].items():
-                            parts.append(f"\n    {def_name}:")
-                            if "properties" in def_info:
-                                def_required = def_info.get("required", [])
-                                for prop_name, prop_info in def_info["properties"].items():
-                                    req_marker = " (required)" if prop_name in def_required else ""
-                                    prop_type = prop_info.get("_annotation", prop_info.get("type", "Unknown"))
-                                    prop_desc = prop_info.get("description", "No description available.")
-                                    default = f" (default: {prop_info['default']})" if "default" in prop_info else ""
-                                    parts.append(
-                                        f"      {prop_name}: {prop_type}{req_marker}{default}\n        {prop_desc}"
-                                    )
-            else:
-                # Simple output schema
-                desc = output_info.get("description", "No description available.")
-                output_type = output_info.get("type", "Unknown")
-                parts.append(f"  Type: {output_type}\n  {desc}")
-        else:
-            parts.append("  Output schema not available")
+        parts.append(format_io_info(io_info))
     except Exception as e:
         parts.append(f"\nFailed to fetch input/output schema: {str(e)}")
 
     return "\n".join(parts)
 
 
+async def search_component_definition(
+    client: AsyncClientProtocol, query: str, model: ModelProtocol, top_k: int = 5
+) -> str:
+    """Searches for components 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
+        top_k: Maximum number of results to return (default: 5)
+
+    Returns:
+        A formatted string containing the matched component definitions
+    """
+    haystack_service = client.haystack_service()
+
+    try:
+        response = await haystack_service.get_component_schemas()
+    except UnexpectedAPIError as e:
+        return f"Failed to retrieve component schemas: {e}"
+
+    components = response["component_schema"]["definitions"]["Components"]
+
+    # Extract text for embedding from all components
+    component_texts: list[tuple[str, str]] = [extract_component_texts(comp) for comp in components.values()]
+    component_types: list[str] = [c[0] for c in component_texts]
+
+    if not component_texts:
+        return "No components found"
+
+    # Compute embeddings
+    query_embedding = model.encode(query)
+    component_embeddings = model.encode([text for _, text in component_texts])
+
+    query_embedding_reshaped = query_embedding.reshape(1, -1)
+
+    # Calculate dot product between target and all paths
+    # This gives us a similarity score for each path
+    similarities = np.dot(component_embeddings, query_embedding_reshaped.T).flatten()
+
+    # Create (path, similarity) pairs
+    component_similarities = list(zip(component_types, similarities, strict=False))
+
+    # Sort by similarity score in descending order
+    component_similarities.sort(key=lambda x: x[1], reverse=True)
+
+    top_components = component_similarities[:top_k]
+    results = []
+    for component_type, sim in top_components:
+        definition = await get_component_definition(client, component_type)
+        results.append(f"Similarity Score: {sim:.3f}\n{definition}\n{'-' * 80}\n")
+
+    return "\n".join(results)
+
+
 async def list_component_families(client: AsyncClientProtocol) -> str:
     """Lists all Haystack component families that are available on deepset."""
     haystack_service = client.haystack_service()

src/deepset_mcp/tools/model_protocol.py

@@ -0,0 +1,16 @@
+from typing import Any, Protocol
+
+import numpy as np
+
+
+class ModelProtocol(Protocol):
+    """Protocol for static embedding models."""
+
+    def encode(self, sentences: list[str] | str) -> np.ndarray[Any, Any]:
+        """
+        Encodes a single or multiple sentences.
+
+        :param sentences: Single sentence or list of sentences to encode
+        :returns: Numpy array of encoded sentences
+        """
+        ...