diff --git a/examples/langgraph/main.py b/examples/langgraph/main.py index 2482a84..47561f8 100644 --- a/examples/langgraph/main.py +++ b/examples/langgraph/main.py @@ -1,7 +1,7 @@ """Use DMA as a retrieval node in a LangGraph StateGraph. -Install the optional framework dependency first: -``pip install 'dma-langgraph[langgraph]'``. +Start the API first with ``make api`` from the repository root, and install the +optional framework dependency: ``pip install 'dma-langgraph[langgraph]'``. """ from __future__ import annotations @@ -20,22 +20,29 @@ class AgentState(TypedDict): response: str -memory = DMAClient(api_key=os.environ["DMA_API_KEY"], agent_id="langgraph-agent") -adapter = DMAMemoryAdapter(memory=memory, query_builder=lambda state: state["user_input"]) - - def respond(state: AgentState) -> dict[str, str]: # Replace this deterministic placeholder with your model invocation. return {"response": f"Relevant memory:\n{state['dma_context']}"} -builder = StateGraph(AgentState) -builder.add_node("recall_memory", adapter.recall_node) -builder.add_node("respond", respond) -builder.add_edge(START, "recall_memory") -builder.add_edge("recall_memory", "respond") -builder.add_edge("respond", END) -graph = builder.compile() +def main() -> None: + with DMAClient( + api_key=os.environ["DMA_API_KEY"], + agent_id="langgraph-agent", + base_url=os.getenv("DMA_BASE_URL", "http://127.0.0.1:8000"), + ) as memory: + adapter = DMAMemoryAdapter(memory=memory, query_builder=lambda state: state["user_input"]) + builder = StateGraph(AgentState) + builder.add_node("recall_memory", adapter.recall_node) + builder.add_node("respond", respond) + builder.add_edge(START, "recall_memory") + builder.add_edge("recall_memory", "respond") + builder.add_edge("respond", END) + graph = builder.compile() + result = graph.invoke({"user_input": "What backend does the user prefer?"}) + + print(result["response"]) + -result = graph.invoke({"user_input": "What backend does the user prefer?"}) -print(result["response"]) +if __name__ == "__main__": + main() diff --git a/packages/dma-mcp/src/dma_mcp/server.py b/packages/dma-mcp/src/dma_mcp/server.py index 44aee83..76f0101 100644 --- a/packages/dma-mcp/src/dma_mcp/server.py +++ b/packages/dma-mcp/src/dma_mcp/server.py @@ -1,6 +1,8 @@ from __future__ import annotations +import logging import os +import sys from dma import DMAClient @@ -35,5 +37,9 @@ def main() -> None: api_key = os.environ.get("DMA_API_KEY") if not api_key: raise SystemExit("DMA_API_KEY must be set") + logging.basicConfig(stream=sys.stderr, level=os.getenv("DMA_MCP_LOG_LEVEL", "INFO")) client = DMAClient(api_key=api_key, agent_id=os.getenv("DMA_MCP_AGENT_ID", "mcp-agent"), base_url=os.getenv("DMA_BASE_URL", "http://127.0.0.1:8000")) - create_server(DMATools(client)).run(transport="stdio") + try: + create_server(DMATools(client)).run(transport="stdio") + finally: + client.close() diff --git a/packages/dma-sdk-python/src/dma/client.py b/packages/dma-sdk-python/src/dma/client.py index ca6c4ce..305f88b 100644 --- a/packages/dma-sdk-python/src/dma/client.py +++ b/packages/dma-sdk-python/src/dma/client.py @@ -2,6 +2,8 @@ from __future__ import annotations +import random +import time from collections.abc import Mapping from datetime import datetime from typing import Any, Self @@ -20,6 +22,8 @@ ) _DEFAULT_BASE_URL = "https://api.dma.dev" +_DEFAULT_MAX_RETRIES = 2 +_RETRY_BACKOFF_SECONDS = 0.2 class DMAClient: @@ -27,6 +31,10 @@ class DMAClient: The caller owns the client lifecycle. Use it as a context manager in scripts and services to ensure its HTTP connection pool is closed. + + Transport failures are retried up to ``max_retries`` times for requests that + are safe to replay: reads, and writes carrying an ``Idempotency-Key``. + Deletes are never retried. """ def __init__( @@ -36,6 +44,7 @@ def __init__( agent_id: str, base_url: str = _DEFAULT_BASE_URL, timeout: float = 5.0, + max_retries: int = _DEFAULT_MAX_RETRIES, transport: httpx.BaseTransport | None = None, ) -> None: if not api_key.strip(): @@ -44,7 +53,10 @@ def __init__( raise ValidationError("agent_id must not be blank") if timeout <= 0: raise ValidationError("timeout must be greater than zero") + if max_retries < 0: + raise ValidationError("max_retries must not be negative") self._agent_id = agent_id + self._max_retries = max_retries self._client = httpx.Client( base_url=base_url.rstrip("/") + "/", timeout=timeout, @@ -75,7 +87,9 @@ def remember( } if expires_at is not None: payload["expires_at"] = expires_at.isoformat() - response = self._request("POST", "v1/memories", json=payload, headers={"Idempotency-Key": key}) + response = self._request( + "POST", "v1/memories", json=payload, headers={"Idempotency-Key": key}, retryable=True + ) return _memory_from_payload(response.json()) def recall( @@ -92,7 +106,7 @@ def recall( payload: dict[str, Any] = {"agent_id": self._agent_id, "query": query, "limit": limit} if types is not None: payload["types"] = [self._memory_type(memory_type).value for memory_type in types] - response = self._request("POST", "v1/memories/recall", json=payload) + response = self._request("POST", "v1/memories/recall", json=payload, retryable=True) return [_recall_result_from_payload(item) for item in response.json()["results"]] def list( @@ -110,7 +124,7 @@ def list( params["type"] = self._memory_type(type).value if cursor is not None: params["cursor"] = cursor - response = self._request("GET", "v1/memories", params=params) + response = self._request("GET", "v1/memories", params=params, retryable=True) payload = response.json() return MemoryPage( items=[_memory_from_payload(item) for item in payload["items"]], @@ -130,7 +144,9 @@ def explain(self, memory_id: str, *, query: str | None = None) -> MemoryExplanat params: dict[str, Any] = {"agent_id": self._agent_id} if query is not None: params["query"] = query - response = self._request("GET", f"v1/memories/{memory_id}/explanation", params=params) + response = self._request( + "GET", f"v1/memories/{memory_id}/explanation", params=params, retryable=True + ) payload = response.json() retrieval = payload["retrieval"] return MemoryExplanation( @@ -155,11 +171,8 @@ def __enter__(self) -> Self: def __exit__(self, *_: object) -> None: self.close() - def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response: - try: - response = self._client.request(method, path, **kwargs) - except httpx.HTTPError as error: - raise DMAConnectionError("unable to reach the DMA API") from error + def _request(self, method: str, path: str, *, retryable: bool = False, **kwargs: Any) -> httpx.Response: + response = self._send(method, path, self._max_retries if retryable else 0, kwargs) if response.status_code == 401: raise AuthenticationError(401, "DMA API key was rejected") if response.is_error: @@ -167,6 +180,16 @@ def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response: raise DMAApiError(response.status_code, message, code=code) return response + def _send(self, method: str, path: str, retries: int, kwargs: dict[str, Any]) -> httpx.Response: + for attempt in range(retries + 1): + try: + return self._client.request(method, path, **kwargs) + except httpx.HTTPError as error: + if attempt == retries: + raise DMAConnectionError("unable to reach the DMA API") from error + time.sleep(_RETRY_BACKOFF_SECONDS * 2**attempt * (0.5 + random.random())) + raise DMAConnectionError("unable to reach the DMA API") + @staticmethod def _validate_non_blank(value: str, field: str) -> None: if not value or not value.strip(): diff --git a/packages/dma-sdk-python/tests/test_client.py b/packages/dma-sdk-python/tests/test_client.py index 9f0558f..944a168 100644 --- a/packages/dma-sdk-python/tests/test_client.py +++ b/packages/dma-sdk-python/tests/test_client.py @@ -5,7 +5,14 @@ import httpx import pytest -from dma import AuthenticationError, DMAApiError, DMAClient, MemoryType, ValidationError +from dma import ( + AuthenticationError, + DMAApiError, + DMAClient, + DMAConnectionError, + MemoryType, + ValidationError, +) def _memory_payload() -> dict[str, object]: @@ -90,6 +97,51 @@ def test_client_surfaces_auth_and_api_errors() -> None: api_client.forget("mem_missing") +def test_retries_replay_safe_requests_with_the_same_idempotency_key() -> None: + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers["Idempotency-Key"]) + if len(seen) == 1: + raise httpx.ConnectError("boom") + return httpx.Response(201, json=_memory_payload()) + + client = DMAClient( + api_key="test-key", + agent_id="coding-agent", + base_url="https://dma.test", + max_retries=1, + transport=httpx.MockTransport(handler), + ) + memory = client.remember(content="User prefers Java Spring Boot.", type=MemoryType.SEMANTIC) + assert memory.id == "mem_abc123" + assert seen == [seen[0], seen[0]] + + +def test_delete_is_not_retried_and_retries_are_bounded() -> None: + attempts: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + attempts.append(request.method) + raise httpx.ConnectError("boom") + + client = DMAClient( + api_key="test-key", + agent_id="coding-agent", + base_url="https://dma.test", + max_retries=2, + transport=httpx.MockTransport(handler), + ) + with pytest.raises(DMAConnectionError): + client.forget("mem_abc123") + assert attempts == ["DELETE"] + + attempts.clear() + with pytest.raises(DMAConnectionError): + client.recall(query="Java") + assert attempts == ["POST", "POST", "POST"] + + def test_client_validates_inputs_before_requests() -> None: with pytest.raises(ValidationError, match="api_key"): DMAClient(api_key=" ", agent_id="coding-agent")