feat: add search functionality to PipelineResource (#92)

* feat: add imports for search models

* feat: add search models for pipeline search functionality

* feat: add imports for search functionality

* feat: add search method with streaming support

* feat: update imports to include search models

* feat: add search method to PipelineResourceProtocol

* feat: add stream field to TransportResponse for streaming support

* feat: add streaming support to AsyncTransport

* feat: add stream parameter to TransportProtocol

* feat: add stream parameter to AsyncClientProtocol

* feat: add imports for search models in tests

* feat: add comprehensive unit tests for search functionality

* feat: add integration test for search functionality

* feat: add stream parameter to BaseFakeClient

* feat: include stream parameter in request recording

* feat: add stream parameter to AsyncDeepsetClient

* feat: pass stream parameter through to transport

* feat: add unit tests for streaming search functionality

* feat: create quick test script to validate search functionality

* chore: remove temporary test file

* fix: update imports to use builtin types instead of typing module

* fix: use builtin types and pipe union syntax in Answer model

* fix: use builtin types and pipe union syntax in Document model

* fix: use builtin types and pipe union syntax in SearchResult model

* fix: use builtin types and pipe union syntax in SearchResponse model

* fix: use pipe union syntax in FilterCondition model

* fix: use builtin types in SearchFilters model

* fix: use builtin types and pipe union syntax in StreamDelta model

* fix: use pipe union syntax in StreamEvent model

* fix: update imports to use builtin types instead of typing module

* feat: add search and search stream

Author: mathislucka

src/deepset_mcp/api/client.py

@@ -1,4 +1,6 @@
 import os
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from types import TracebackType
 from typing import Any, Self, TypeVar, overload
 
@@ -8,7 +10,12 @@
 from deepset_mcp.api.pipeline.resource import PipelineResource
 from deepset_mcp.api.pipeline_template.resource import PipelineTemplateResource
 from deepset_mcp.api.protocols import AsyncClientProtocol
-from deepset_mcp.api.transport import AsyncTransport, TransportProtocol, TransportResponse
+from deepset_mcp.api.transport import (
+    AsyncTransport,
+    StreamingResponse,
+    TransportProtocol,
+    TransportResponse,
+)
 from deepset_mcp.api.user.resource import UserResource
 
 T = TypeVar("T")
@@ -85,7 +92,29 @@ async def request(
         response_type: type[T] | None = None,
         **kwargs: Any,
     ) -> TransportResponse[Any]:
-        """Make a request to the deepset API."""
+        """
+        Make a regular (non-streaming) request to the deepset API.
+
+        Parameters
+        ----------
+        endpoint : str
+            API endpoint path
+        method : str, default="GET"
+            HTTP method
+        data : dict, optional
+            JSON data to send in request body
+        headers : dict, optional
+            Additional headers to include
+        response_type : type[T], optional
+            Expected response type for type checking
+        **kwargs : Any
+            Additional arguments to pass to transport
+
+        Returns
+        -------
+        TransportResponse[T]
+            Response with parsed JSON if available
+        """
         if not endpoint.startswith("/"):
             endpoint = f"/{endpoint}"
         url = self.base_url + endpoint
@@ -111,6 +140,81 @@ async def request(
             **kwargs,
         )
 
+    def stream_request(
+        self,
+        endpoint: str,
+        *,
+        method: str = "POST",
+        data: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> AbstractAsyncContextManager[StreamingResponse]:
+        """
+        Make a streaming request to the deepset API.
+
+        Must be used as an async context manager to ensure proper cleanup.
+
+        Parameters
+        ----------
+        endpoint : str
+            API endpoint path
+        method : str, default="POST"
+            HTTP method (usually POST for streaming)
+        data : dict, optional
+            JSON data to send in request body
+        headers : dict, optional
+            Additional headers to include
+        **kwargs : Any
+            Additional arguments to pass to transport
+
+        Yields
+        ------
+        StreamingResponse
+            Response object with streaming capabilities
+
+        Examples
+        --------
+        async with client.stream_request("/pipelines/search-stream", data={"query": "AI"}) as response:
+            if response.success:
+                async for line in response.iter_lines():
+                    # Process each line of the stream
+                    data = json.loads(line)
+                    print(data)
+            else:
+                # Handle error
+                error_body = await response.read_body()
+                print(f"Error {response.status_code}: {error_body}")
+        """
+
+        @asynccontextmanager
+        async def _stream() -> AsyncIterator[StreamingResponse]:
+            if not endpoint.startswith("/"):
+                full_endpoint = f"/{endpoint}"
+            url = self.base_url + full_endpoint
+
+            # Default headers for streaming
+            request_headers: dict[str, str] = {
+                "Authorization": f"Bearer {self.api_key}",
+                "Accept": "text/event-stream,application/json,text/plain,*/*",
+            }
+            if data is not None:
+                request_headers["Content-Type"] = "application/json"
+            # Merge custom headers
+            if headers:
+                headers.setdefault("Authorization", request_headers["Authorization"])
+                request_headers.update(headers)
+
+            async with self._transport.stream(
+                method,
+                url,
+                json=data,
+                headers=request_headers,
+                **kwargs,
+            ) as response:
+                yield response
+
+        return _stream()
+
     async def close(self) -> None:
         """Close underlying transport resources."""
         await self._transport.close()

src/deepset_mcp/api/pipeline/models.py

@@ -1,8 +1,9 @@
 from datetime import datetime
 from enum import StrEnum
 from typing import Any
+from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from deepset_mcp.api.shared_models import DeepsetUser
 
@@ -96,3 +97,101 @@ class PipelineLogList(BaseModel):
     data: list[PipelineLog]
     has_more: bool
     total: int
+
+
+# Search-related models
+
+
+class OffsetRange(BaseModel):
+    """Model representing an offset range."""
+
+    start: int
+    end: int
+
+
+class DeepsetAnswer(BaseModel):
+    """Model representing a search answer."""
+
+    answer: str  # Required field
+    context: str | None = None
+    document_id: str | None = None
+    document_ids: list[str] | None = None
+    file: dict[str, Any] | None = None
+    files: list[dict[str, Any]] | None = None
+    meta: dict[str, Any] | None = None
+    offsets_in_context: list[OffsetRange] | None = None
+    offsets_in_document: list[OffsetRange] | None = None
+    prompt: str | None = None
+    result_id: UUID | None = None
+    score: float | None = None
+    type: str | None = None
+
+
+class DeepsetDocument(BaseModel):
+    """Model representing a search document."""
+
+    content: str  # Required field
+    meta: dict[str, Any]  # Required field - can hold any value
+    embedding: list[float] | None = None
+    file: dict[str, Any] | None = None
+    id: str | None = None
+    result_id: UUID | None = None
+    score: float | None = None
+
+
+class DeepsetSearchResponse(BaseModel):
+    """Model representing a single search result."""
+
+    debug: dict[str, Any] | None = Field(default=None, alias="_debug")
+    answers: list[DeepsetAnswer] = Field(default_factory=list)
+    documents: list[DeepsetDocument] = Field(default_factory=list)
+    prompts: dict[str, str] | None = None
+    query: str | None = None
+    query_id: UUID | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_response(cls, data: dict[str, Any]) -> dict[str, Any]:
+        """Normalize the response from the search and search-stream endpoints.
+
+        The search endpoint returns a list of results, but we only ever use the first result.
+        We are not sending batch queries, so there will never be more than one result.
+        We use this validator to transform the data so that we can use the same response model for search and
+            search-stream endpoints.
+        """
+        # Handle non-stream format with 'results' array
+        if "results" in data and isinstance(data["results"], list):
+            if len(data["results"]) > 0:
+                first_result = data["results"][
+                    0
+                ]  # we only ever care for the first result as we don't use batch queries
+                normalized = {
+                    "query_id": data.get("query_id", first_result.get("query_id")),
+                    "query": first_result.get("query"),
+                    "answers": first_result.get("answers", []),
+                    "documents": first_result.get("documents", []),
+                    "prompts": first_result.get("prompts"),
+                    "_debug": first_result.get("_debug") or first_result.get("debug"),
+                }
+                return normalized
+            else:
+                return {}
+        else:
+            return data
+
+
+class StreamDelta(BaseModel):
+    """Model representing a streaming delta."""
+
+    text: str
+    meta: dict[str, Any] | None = None
+
+
+class DeepsetStreamEvent(BaseModel):
+    """Model representing a stream event."""
+
+    query_id: str | UUID | None = None
+    type: str  # "delta", "result", or "error"
+    delta: StreamDelta | None = None
+    result: DeepsetSearchResponse | None = None
+    error: str | None = None

src/deepset_mcp/api/pipeline/resource.py

@@ -1,10 +1,14 @@
+import json
 import logging
+from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
 
 from deepset_mcp.api.exceptions import UnexpectedAPIError
 from deepset_mcp.api.pipeline.log_level import LogLevel
 from deepset_mcp.api.pipeline.models import (
     DeepsetPipeline,
+    DeepsetSearchResponse,
+    DeepsetStreamEvent,
     NoContentResponse,
     PipelineLogList,
     PipelineValidationResult,
@@ -261,3 +265,103 @@ async def delete(self, pipeline_name: str) -> NoContentResponse:
         raise_for_status(resp)
 
         return NoContentResponse(message="Pipeline deleted successfully.")
+
+    async def search(
+        self,
+        pipeline_name: str,
+        query: str,
+        debug: bool = False,
+        view_prompts: bool = False,
+        params: dict[str, Any] | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> DeepsetSearchResponse:
+        """Search using a pipeline.
+
+        :param pipeline_name: Name of the pipeline to use for search.
+        :param query: Search query.
+        :param debug: Whether to include debug information.
+        :param view_prompts: Whether to include prompts in the response.
+        :param params: Additional parameters for pipeline components.
+        :param filters: Search filters to apply.
+
+        :returns: SearchResponse containing search results.
+        """
+        # Prepare request data
+        data: dict[str, Any] = {
+            "queries": [query],  # API expects a list but we only send one query
+            "debug": debug,
+            "view_prompts": view_prompts,
+        }
+
+        if params:
+            data["params"] = params
+
+        if filters:
+            data["filters"] = filters
+
+        resp = await self._client.request(
+            endpoint=f"v1/workspaces/{self._workspace}/pipelines/{pipeline_name}/search",
+            method="POST",
+            data=data,
+            response_type=dict[str, Any],
+        )
+
+        raise_for_status(resp)
+
+        if resp.json is not None:
+            return DeepsetSearchResponse.model_validate(resp.json)
+        else:
+            # Return empty response if no JSON data
+            return DeepsetSearchResponse()
+
+    async def search_stream(
+        self,
+        pipeline_name: str,
+        query: str,
+        debug: bool = False,
+        view_prompts: bool = False,
+        params: dict[str, Any] | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> AsyncIterator[DeepsetStreamEvent]:
+        """Search using a pipeline with response streaming.
+
+        :param pipeline_name: Name of the pipeline to use for search.
+        :param query: Search query.
+        :param debug: Whether to include debug information.
+        :param view_prompts: Whether to include prompts in the response.
+        :param params: Additional parameters for pipeline components.
+        :param filters: Search filters to apply.
+
+        :returns: AsyncIterator streaming the result.
+        """
+        # For streaming, we need to add include_result flag
+        # Prepare request data
+        data: dict[str, Any] = {
+            "query": query,
+            "debug": debug,
+            "view_prompts": view_prompts,
+            "include_result": True,
+        }
+
+        if params:
+            data["params"] = params
+
+        if filters:
+            data["filters"] = filters
+
+        async with self._client.stream_request(
+            endpoint=f"v1/workspaces/{self._workspace}/pipelines/{pipeline_name}/search-stream",
+            method="POST",
+            data=data,
+        ) as resp:
+            async for line in resp.iter_lines():
+                try:
+                    event_dict = json.loads(line)
+                    event = DeepsetStreamEvent.model_validate(event_dict)
+
+                    if event.error is not None:
+                        raise UnexpectedAPIError(message=event.error)
+                    yield event
+                except (json.JSONDecodeError, ValueError):
+                    # Skip malformed events
+                    continue