Merge pull request #12 from deepset-ai/feat/deepset_client

feat: add AsyncDeepsetClient

Author: mathislucka

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "mcp[cli]>=1.6.0",
     "requests>=2.32.3",
     "uvicorn>=0.34.2",
+    "httpx",
 ]
 
 [project.scripts]
@@ -23,7 +24,8 @@ packages = ["src/deepset_mcp"]
 
 [dependency-groups]
 dev = [
-  "pytest"
+    "pytest",
+    "pytest-asyncio",
 ]
 lint = [
     "ruff",
@@ -60,6 +62,9 @@ ignore = [
 ]
 isort = { combine-as-imports = true, known-first-party = ["deepset_mcp"] }
 
+[tool.ruff.lint.per-file-ignores]
+"test/*" = ["D"]
+
 [tool.ruff.lint.pydocstyle]
 convention = "pep257"
 

src/deepset_mcp/client.py

@@ -0,0 +1,89 @@
+import os
+from types import TracebackType
+from typing import Any
+
+from deepset_mcp.transport import AsyncTransport, TransportProtocol
+
+
+class AsyncDeepsetClient:
+    """Async Client for interacting with the deepset API."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str = "https://api.cloud.deepset.ai/api",
+        transport: TransportProtocol | None = None,
+        transport_config: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Initialize an instance of the AsyncDeepsetClient.
+
+        Parameters
+        ----------
+        api_key : str, optional
+            API key or token. Falls back to DEEPSET_API_KEY env var.
+        base_url : str, optional
+            Base URL for the deepset API.
+        transport : TransportProtocol, optional
+            Custom transport implementation.
+        transport_config : dict, optional
+            Configuration for default transport (e.g. timeout).
+        """
+        self.api_key = api_key or os.environ.get("DEEPSET_API_KEY")
+        if not self.api_key:
+            raise ValueError("API key not provided and DEEPSET_API_KEY environment variable not set")
+        self.base_url = base_url
+        if transport is not None:
+            self._transport = transport
+        else:
+            self._transport = AsyncTransport(
+                base_url=self.base_url,
+                api_key=self.api_key,
+                config=transport_config,
+            )
+
+    async def request(
+        self,
+        endpoint: str,
+        method: str = "GET",
+        data: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Any:
+        """Make a request to the deepset API."""
+        if not endpoint.startswith("/"):
+            endpoint = f"/{endpoint}"
+        url = self.base_url + endpoint
+
+        # Default headers
+        request_headers: dict[str, str] = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Accept": "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)
+
+        return await self._transport.request(
+            method,
+            url,
+            json=data,
+            headers=request_headers,
+        )
+
+    async def close(self) -> None:
+        """Close underlying transport resources."""
+        await self._transport.close()
+
+    async def __aenter__(self) -> "AsyncDeepsetClient":
+        """Enter the AsyncContextManager."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> bool:
+        """Exit the AsyncContextmanager and clean up resources."""
+        await self.close()
+        return False

src/deepset_mcp/main.py

@@ -5,6 +5,8 @@
 from mcp.server.fastmcp import FastMCP
 from requests import HTTPError
 
+from deepset_mcp.client import AsyncDeepsetClient
+
 # Initialize MCP Server
 mcp = FastMCP("Deepset Cloud MCP")
 
@@ -69,20 +71,18 @@ def deepset_api_request(endpoint: str, method: str = "GET", data: dict[str, Any]
 
 
 @mcp.tool()
-def list_pipelines() -> dict[str, Any]:
+async def list_pipelines() -> Any:
     """Retrieves a list of all pipelines available within the currently configured deepset workspace.
 
     Use this when you need to know the names or IDs of existing pipelines.
     """
     workspace = get_workspace()
-    try:
-        return deepset_api_request(f"/workspaces/{workspace}/pipelines")
-    except Exception as e:
-        return {"error": str(e)}
+    async with AsyncDeepsetClient() as client:
+        return await client.request(endpoint=f"v1/workspaces/{workspace}/pipelines", method="GET")
 
 
 @mcp.tool()
-def get_pipeline(pipeline_id: str) -> dict[str, Any]:
+async def get_pipeline(pipeline_id: str) -> Any:
     """Fetches detailed configuration information for a specific pipeline, identified by its unique `pipeline_id`.
 
     This includes its components, connections, and metadata.
@@ -91,23 +91,19 @@ def get_pipeline(pipeline_id: str) -> dict[str, Any]:
     :param pipeline_id: ID of the pipeline to retrieve.
     """
     workspace = get_workspace()
-    try:
-        return deepset_api_request(f"/workspaces/{workspace}/pipelines/{pipeline_id}")
-    except Exception as e:
-        return {"error": str(e)}
+    async with AsyncDeepsetClient() as client:
+        return await client.request(endpoint=f"v1/workspaces/{workspace}/pipelines/{pipeline_id}", method="GET")
 
 
 @mcp.tool()
-def get_component_schemas() -> dict[str, Any]:
+async def get_component_schemas() -> Any:
     """Retrieves the schemas for all available Haystack components from the deepset API.
 
     These schemas define the expected input and output parameters for each component type, which is useful for
     constructing or validating componets in a pipeline YAML.
     """
-    try:
-        return deepset_api_request("/haystack/components")
-    except Exception as e:
-        return {"error": str(e)}
+    async with AsyncDeepsetClient() as client:
+        return await client.request(endpoint="v1/haystack/components", method="GET")
 
 
 @mcp.tool()

src/deepset_mcp/transport.py

@@ -0,0 +1,60 @@
+from typing import Any, Protocol
+
+import httpx
+
+
+class TransportProtocol(Protocol):
+    """Protocol for HTTP transport."""
+
+    async def request(self, method: str, url: str, **kwargs: Any) -> Any:
+        """Send an HTTP request and return the parsed JSON response."""
+        ...
+
+    async def close(self) -> None:
+        """Clean up any resources (e.g., close connections)."""
+        ...
+
+
+class AsyncTransport(TransportProtocol):
+    """Asynchronous HTTP transport using httpx.AsyncClient."""
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str,
+        config: dict[str, Any] | None = None,
+    ):
+        """
+        Initialize an instance of AsyncTransport.
+
+        Parameters
+        ----------
+        base_url : str
+            Base URL for the API
+        api_key : str
+            Bearer token for authentication
+        config : dict, optional
+            Configuration for httpx.AsyncClient, e.g., {'timeout': 10.0}
+        """
+        config = config or {}
+        # Ensure auth header
+        headers = config.pop("headers", {})
+        headers.setdefault("Authorization", f"Bearer {api_key}")
+        # Build client kwargs
+        client_kwargs = {
+            "base_url": base_url,
+            "headers": headers,
+            **config,
+        }
+        self._client = httpx.AsyncClient(**client_kwargs)
+
+    async def request(self, method: str, url: str, **kwargs: Any) -> Any:
+        """Send an HTTP request and return the parsed JSON response."""
+        resp = await self._client.request(method, url, **kwargs)
+        resp.raise_for_status()
+
+        return resp.json()
+
+    async def close(self) -> None:
+        """Clean up any resources (e.g., close connections)."""
+        await self._client.aclose()

test/test_async_deepset_client.py

@@ -0,0 +1,112 @@
+from typing import Any
+
+import pytest
+import pytest_asyncio
+
+from deepset_mcp.client import AsyncDeepsetClient
+from deepset_mcp.transport import AsyncTransport, TransportProtocol
+
+
+class DummyProtocol(TransportProtocol):
+    def __init__(self) -> None:
+        self.requests: list[dict[str, Any]] = []
+        self.closed: bool = False
+
+    async def request(self, method: str, url: str, **kwargs: Any) -> Any:
+        # Record the request and return a dummy response
+        record: dict[str, Any] = {"method": method, "url": url, **kwargs}
+        self.requests.append(record)
+        return {"dummy": "response"}
+
+    async def close(self) -> None:
+        self.closed = True
+
+
+@pytest_asyncio.fixture(autouse=True)  # type: ignore
+def clear_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Ensure DEEPSET_API_KEY is unset by default unless explicitly set
+    monkeypatch.delenv("DEEPSET_API_KEY", raising=False)
+
+
+@pytest.mark.asyncio
+async def test_init_with_api_key_and_transport() -> None:
+    dummy: DummyProtocol = DummyProtocol()
+    client: AsyncDeepsetClient = AsyncDeepsetClient(api_key="testkey", transport=dummy)
+    assert client.api_key == "testkey"
+    assert client._transport is dummy
+
+
+@pytest.mark.asyncio
+async def test_init_without_api_key_raises() -> None:
+    # No API key in args or env
+    with pytest.raises(ValueError):
+        AsyncDeepsetClient()
+
+
+@pytest.mark.asyncio
+async def test_init_with_env_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("DEEPSET_API_KEY", "envkey")
+    client: AsyncDeepsetClient = AsyncDeepsetClient()
+    assert client.api_key == "envkey"
+    # Default transport is AsyncTransport
+    assert isinstance(client._transport, AsyncTransport)
+
+
+@pytest.mark.asyncio
+async def test_request_default_headers_and_url() -> None:
+    dummy: DummyProtocol = DummyProtocol()
+    client: AsyncDeepsetClient = AsyncDeepsetClient(api_key="key", base_url="https://api.test", transport=dummy)
+
+    resp: Any = await client.request("endpoint")
+    # Should normalize endpoint to '/endpoint'
+    assert resp == {"dummy": "response"}
+    assert len(dummy.requests) == 1
+
+    call: dict[str, Any] = dummy.requests[0]
+    assert call["method"] == "GET"
+    assert call["url"] == "https://api.test/endpoint"
+    headers: dict[str, str] = call["headers"]
+    assert headers["Authorization"] == "Bearer key"
+    assert headers["Accept"] == "application/json,text/plain,*/*"
+    assert "Content-Type" not in headers
+    assert call.get("json") is None
+
+
+@pytest.mark.asyncio
+async def test_request_with_data_and_custom_headers() -> None:
+    dummy: DummyProtocol = DummyProtocol()
+    client: AsyncDeepsetClient = AsyncDeepsetClient(api_key="key", base_url="https://api.test", transport=dummy)
+
+    data: dict[str, Any] = {"foo": "bar"}
+    custom: dict[str, str] = {"X-Custom": "value"}
+    resp: Any = await client.request("/path", method="POST", data=data, headers=custom)
+    assert resp == {"dummy": "response"}
+    assert len(dummy.requests) == 1
+
+    call = dummy.requests[0]
+    assert call["method"] == "POST"
+    assert call["url"] == "https://api.test/path"
+    headers = call["headers"]
+    # Custom header merged
+    assert headers["X-Custom"] == "value"
+    assert headers["Content-Type"] == "application/json"
+    # Authorization preserved
+    assert headers["Authorization"] == "Bearer key"
+    assert call.get("json") == data
+
+
+@pytest.mark.asyncio
+async def test_close_and_context_manager() -> None:
+    dummy: DummyProtocol = DummyProtocol()
+    client: AsyncDeepsetClient = AsyncDeepsetClient(api_key="key", transport=dummy)
+
+    # Test close
+    await client.close()
+    assert dummy.closed is True
+
+    # Test async context manager
+    dummy2: DummyProtocol = DummyProtocol()
+    async with AsyncDeepsetClient(api_key="key", transport=dummy2) as ctx:
+        assert isinstance(ctx, AsyncDeepsetClient)
+    # After exit, close should be called
+    assert dummy2.closed is True