feat: template tools should support indexes (#153)

* feat: template tools should support indexes

* fix: format

Author: mathislucka

src/deepset_mcp/api/pipeline_template/models.py

@@ -1,7 +1,8 @@
 from enum import StrEnum
+from typing import Any
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 
 class PipelineType(StrEnum):
@@ -28,10 +29,36 @@ class PipelineTemplate(BaseModel):
     display_name: str = Field(alias="name")
     pipeline_template_id: UUID = Field(alias="pipeline_template_id")
     potential_applications: list[str] = Field(alias="potential_applications")
-    yaml_config: str | None = Field(None, alias="query_yaml")
+    yaml_config: str | None = None
     tags: list[PipelineTemplateTag]
     pipeline_type: PipelineType
 
+    @model_validator(mode="before")
+    @classmethod
+    def populate_yaml_config(cls, values: Any) -> Any:
+        """Populate yaml_config from query_yaml or indexing_yaml based on pipeline_type."""
+        if not isinstance(values, dict):
+            return values
+
+        # Skip if yaml_config is already set
+        if values.get("yaml_config") is not None:
+            return values
+
+        # Get pipeline_type from the model data
+        pipeline_type = values.get("pipeline_type")
+
+        if pipeline_type == PipelineType.INDEXING or pipeline_type == "indexing":
+            yaml_config = values.get("indexing_yaml")
+        elif pipeline_type == PipelineType.QUERY or pipeline_type == "query":
+            yaml_config = values.get("query_yaml")
+        else:
+            yaml_config = None
+
+        if yaml_config is not None:
+            values["yaml_config"] = yaml_config
+
+        return values
+
 
 class PipelineTemplateList(BaseModel):
     """Response model for listing pipeline templates."""

src/deepset_mcp/tool_factory.py

@@ -47,9 +47,9 @@
     validate_pipeline as validate_pipeline_tool,
 )
 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,
+    get_template as get_pipeline_template_tool,
+    list_templates as list_pipeline_templates_tool,
+    search_templates as search_pipeline_templates_tool,
 )
 from deepset_mcp.tools.secrets import (
     get_secret as get_secret_tool,
@@ -230,20 +230,20 @@ def get_workspace_from_env() -> str:
         deploy_index_tool,
         ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
     ),
-    "list_pipeline_templates": (
+    "list_templates": (
         list_pipeline_templates_tool,
         ToolConfig(
             needs_client=True,
             needs_workspace=True,
             memory_type=MemoryType.EXPLORABLE,
-            custom_args={"filter": None, "field": "created_at", "order": "DESC", "limit": 100},
+            custom_args={"field": "created_at", "order": "DESC", "limit": 100},
         ),
     ),
-    "get_pipeline_template": (
+    "get_template": (
         get_pipeline_template_tool,
         ToolConfig(needs_client=True, needs_workspace=True, memory_type=MemoryType.EXPLORABLE),
     ),
-    "search_pipeline_templates": (
+    "search_templates": (
         search_pipeline_templates_tool,
         ToolConfig(
             needs_client=True,

src/deepset_mcp/tools/pipeline_template.py

@@ -6,51 +6,53 @@
     PipelineTemplateList,
     PipelineTemplateSearchResult,
     PipelineTemplateSearchResults,
+    PipelineType,
 )
 from deepset_mcp.api.protocols import AsyncClientProtocol
 from deepset_mcp.tools.model_protocol import ModelProtocol
 
 
-async def list_pipeline_templates(
+async def list_templates(
     *,
     client: AsyncClientProtocol,
     workspace: str,
     limit: int = 100,
     field: str = "created_at",
     order: str = "DESC",
-    filter: str | None = None,
+    pipeline_type: PipelineType | str | None = None,
 ) -> PipelineTemplateList | str:
-    """Retrieves a list of all available pipeline templates.
+    """Retrieves a list of all available pipeline and indexing templates.
 
     :param client: The async client for API requests.
     :param workspace: The workspace to list templates from.
     :param limit: Maximum number of templates to return (default: 100).
     :param field: Field to sort by (default: "created_at").
     :param order: Sort order, either "ASC" or "DESC" (default: "DESC").
-    :param filter: OData filter expression to filter templates by criteria.
+    :param pipeline_type: The type of pipeline to return.
 
     :returns: List of pipeline templates or error message.
     """
     try:
         return await client.pipeline_templates(workspace=workspace).list_templates(
-            limit=limit, field=field, order=order, filter=filter
+            limit=limit,
+            field=field,
+            order=order,
+            filter=f"pipeline_type eq '{pipeline_type}'" if pipeline_type else None,
         )
     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
-) -> PipelineTemplate | str:
-    """Fetches detailed information for a specific pipeline template, identified by its `template_name`.
+async def get_template(*, client: AsyncClientProtocol, workspace: str, template_name: str) -> PipelineTemplate | str:
+    """Fetches detailed information for a specific pipeline or indexing 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.
+    :returns: Pipeline or indexing template details or error message.
     """
     try:
         return await client.pipeline_templates(workspace=workspace).get_template(template_name=template_name)
@@ -60,22 +62,29 @@ async def get_pipeline_template(
         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 = 10
+async def search_templates(
+    *,
+    client: AsyncClientProtocol,
+    query: str,
+    model: ModelProtocol,
+    workspace: str,
+    top_k: int = 10,
+    pipeline_type: PipelineType | str = PipelineType.QUERY,
 ) -> PipelineTemplateSearchResults | str:
-    """Searches for pipeline templates based on name or description using semantic similarity.
+    """Searches for pipeline or indexing templates based on name or description using semantic similarity.
 
     :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).
+    :param pipeline_type: The type of pipeline to return ('indexing' or 'query'; default: 'query').
 
     :returns: Search results with similarity scores or error message.
     """
     try:
         response = await client.pipeline_templates(workspace=workspace).list_templates(
-            filter="pipeline_type eq 'QUERY'"
+            filter=f"pipeline_type eq '{pipeline_type}'"
         )
     except UnexpectedAPIError as e:
         return f"Failed to retrieve pipeline templates: {e}"

test/integration/test_integration_pipeline_template_resource.py

@@ -143,3 +143,86 @@ async def test_list_templates_with_custom_sorting(
     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
+
+
+@pytest.mark.asyncio
+async def test_get_indexing_template(
+    template_resource: PipelineTemplateResource,
+) -> None:
+    """Test getting a single indexing template by name.
+
+    First lists all indexing templates, then gets the first one by name.
+    """
+    # Get all indexing templates
+    indexing_templates_list = await template_resource.list_templates(filter="pipeline_type eq 'INDEXING'")
+
+    # Skip if no indexing templates are available
+    if not indexing_templates_list.data:
+        pytest.skip("No indexing templates available in the test environment")
+
+    # Get the first indexing template's name
+    template_name = indexing_templates_list.data[0].template_name
+
+    # Now get that specific template
+    template = await template_resource.get_template(template_name=template_name)
+
+    # Verify the template was retrieved correctly
+    assert template.template_name == template_name
+    assert template.pipeline_type == "indexing"
+    assert template.pipeline_template_id is not None
+    assert template.yaml_config is not None  # Should have indexing_yaml content
+    assert isinstance(template.best_for, list)
+    assert isinstance(template.potential_applications, list)
+    assert isinstance(template.tags, list)
+
+
+@pytest.mark.asyncio
+async def test_list_indexing_templates_with_filter(
+    template_resource: PipelineTemplateResource,
+) -> None:
+    """Test listing templates with an indexing pipeline type filter."""
+    # Test filtering by INDEXING pipeline type
+    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, PipelineTemplateList)
+    assert isinstance(indexing_templates_list.data, list)
+
+    # If templates are available, verify they are all INDEXING type
+    for template in indexing_templates_list.data:
+        assert isinstance(template, PipelineTemplate)
+        assert template.pipeline_type == "indexing"
+        assert template.yaml_config is not None  # Should have indexing_yaml content
+
+
+@pytest.mark.asyncio
+async def test_mixed_pipeline_types_integration(
+    template_resource: PipelineTemplateResource,
+) -> None:
+    """Test that query and indexing templates can be retrieved together."""
+    # Get all templates
+    all_templates = await template_resource.list_templates()
+
+    # Skip if no templates are available
+    if not all_templates.data:
+        pytest.skip("No templates available in the test environment")
+
+    # Separate templates by type
+    query_templates = [t for t in all_templates.data if t.pipeline_type == "query"]
+    indexing_templates = [t for t in all_templates.data if t.pipeline_type == "indexing"]
+
+    # Verify that both types can exist and have proper yaml_config
+    for template in query_templates:
+        assert template.pipeline_type == "query"
+        if template.yaml_config is not None:
+            # Query templates should have query_yaml content
+            assert isinstance(template.yaml_config, str)
+
+    for template in indexing_templates:
+        assert template.pipeline_type == "indexing"
+        if template.yaml_config is not None:
+            # Indexing templates should have indexing_yaml content
+            assert isinstance(template.yaml_config, str)
+
+    # Verify that the total matches the sum of individual types
+    assert len(query_templates) + len(indexing_templates) == len(all_templates.data)