"""Minimal asynchronous client for the LightRAG APIs DeerFlow consumes.""" from __future__ import annotations from typing import Any import httpx QUERY_MODES = ("naive", "local", "global", "hybrid", "mix") class LightRAGError(Exception): """Base class for normalized LightRAG failures.""" class LightRAGAPIError(LightRAGError): """LightRAG rejected the request with a readable failure.""" class LightRAGConnectionError(LightRAGError): """LightRAG could not be reached or timed out.""" class LightRAGProtocolError(LightRAGError): """LightRAG returned an invalid or unexpected HTTP response.""" class LightRAGClient: """Direct HTTP client for DeerFlow's read-only retrieval tools. The client deliberately owns no cache or persistent state. A fresh HTTP session is opened for each method call so callers do not need to manage a client lifecycle. The optional API key is sent as the ``X-API-Key`` request header, the single credential form LightRAG documents for API-key-authenticated servers; unauthenticated deployments simply omit it. """ def __init__( self, *, base_url: str, api_key: str | None, timeout: float = 30, transport: httpx.AsyncBaseTransport | None = None, ) -> None: self.base_url = base_url.rstrip("/") self.timeout = timeout self._api_key = api_key self._transport = transport def _redact(self, value: object) -> str: text = str(value) if self._api_key: text = text.replace(self._api_key, "[REDACTED]") return text def _error_message(self, payload: object, status_code: int) -> str | None: """Extract a redacted, human-readable message from an error payload. LightRAG failures carry text in ``message`` (QueryDataResponse envelope) or ``detail`` (FastAPI error handler, either a string or a list of validation objects whose ``msg`` holds the reason, prefixed by pydantic's "Value error, "). Anything else — structured bodies, plain text, missing payloads — yields ``None`` so the caller falls back to a stable protocol error instead of dumping raw JSON at the model. """ if not isinstance(payload, dict): return None candidate = payload.get("message") if not isinstance(candidate, str) or not candidate.strip(): detail = payload.get("detail") if isinstance(detail, str): candidate = detail elif isinstance(detail, list): candidate = self._first_validation_message(detail) if isinstance(candidate, str) and candidate.strip(): text = candidate.removeprefix("Value error, ").strip() return self._redact(text) return None @staticmethod def _first_validation_message(items: list[object]) -> str | None: for item in items: if isinstance(item, dict): message = item.get("msg") if isinstance(message, str) and message.strip(): return message return None async def _request(self, method: str, path: str, *, json: dict[str, Any] | None = None) -> dict[str, Any]: request_headers = {"Accept": "application/json"} if self._api_key: request_headers["X-API-Key"] = self._api_key client_kwargs: dict[str, Any] = { "base_url": self.base_url, "headers": request_headers, "timeout": self.timeout, } if self._transport is not None: client_kwargs["transport"] = self._transport try: async with httpx.AsyncClient(**client_kwargs) as client: response = await client.request(method, path, json=json) except httpx.TimeoutException: raise LightRAGConnectionError(f"LightRAG request timed out after {self.timeout:g} seconds.") from None except httpx.RequestError as exc: detail = self._redact(exc) raise LightRAGConnectionError(f"{type(exc).__name__}: {detail}") from None if response.is_error: # A 404 on the data-retrieval endpoint means either a wrong # base_url or a LightRAG older than v1.4.9, where /query/data did # not exist yet; the default "Not Found" body helps neither case. if response.status_code == 404: raise LightRAGAPIError("LightRAG data-retrieval endpoint not found; check base_url or upgrade LightRAG to v1.4.9 or newer.") try: error_payload = response.json() except ValueError: raise LightRAGProtocolError(f"LightRAG request failed (HTTP {response.status_code}).") from None message = self._error_message(error_payload, response.status_code) if message is not None: raise LightRAGAPIError(message) raise LightRAGProtocolError(f"LightRAG request failed (HTTP {response.status_code}).") try: payload = response.json() except ValueError: raise LightRAGProtocolError("LightRAG returned invalid JSON.") from None if not isinstance(payload, dict): raise LightRAGProtocolError("LightRAG returned a non-object JSON payload.") status = payload.get("status") if status != "success": if "chunks" in payload or "entities" in payload: # v1.4.8 answers with a flat {entities, relationships, # chunks, metadata} payload; the status/data envelope and the # reference-bearing chunk fields both shipped in v1.4.9. raise LightRAGAPIError("LightRAG server response predates v1.4.9; upgrade LightRAG to v1.4.9 or newer to use the data-retrieval endpoint.") message = payload.get("message") text = self._redact(message) if isinstance(message, str) and message.strip() else "LightRAG request failed." raise LightRAGAPIError(text) return payload async def query_data( self, query: str, *, mode: str = "hybrid", top_k: int = 60, chunk_top_k: int | None = None, ) -> dict[str, Any]: """Run one read-only structured retrieval against ``POST /query/data``. The data endpoint performs no LLM generation and always returns entities, relationships, chunks, and references, which is exactly the read-only shape DeerFlow's knowledge tool consumes. """ if mode not in QUERY_MODES: raise ValueError(f"mode must be one of {QUERY_MODES}") request_body: dict[str, object] = {"query": query, "mode": mode, "top_k": top_k} if chunk_top_k is not None: request_body["chunk_top_k"] = chunk_top_k payload = await self._request("POST", "/query/data", json=request_body) data = payload.get("data") if not isinstance(data, dict): raise LightRAGProtocolError("LightRAG returned an invalid retrieval result.") return data