knowledge-adapter-service/src/knowledge_adapter/weknora_client.py

from types import TracebackType
from uuid import UUID

import httpx
from pydantic import SecretStr

from knowledge_adapter.observability import current_request_id


class WeKnoraUnavailableError(Exception):
    """WeKnora could not be reached within the configured budget."""


class WeKnoraBadResponseError(Exception):
    """WeKnora returned an unusable response."""


class WeKnoraClient:
    def __init__(
        self,
        *,
        base_url: str,
        api_key: SecretStr,
        timeout_seconds: float,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        short_timeout = min(5.0, timeout_seconds)
        timeout = httpx.Timeout(
            timeout_seconds,
            connect=short_timeout,
            read=timeout_seconds,
            write=short_timeout,
            pool=short_timeout,
        )
        limits = httpx.Limits(max_connections=100, max_keepalive_connections=20)
        self._client = httpx.AsyncClient(
            base_url=base_url.rstrip("/"),
            headers={"X-API-Key": api_key.get_secret_value()},
            timeout=timeout,
            limits=limits,
            transport=transport,
            trust_env=False,
        )

    async def __aenter__(self) -> "WeKnoraClient":
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        await self.aclose()

    async def aclose(self) -> None:
        await self._client.aclose()

    async def search(
        self,
        *,
        query: str,
        knowledge_base_id: UUID,
    ) -> list[dict[str, object]]:
        request_id = current_request_id()
        headers = {"X-Request-ID": request_id} if request_id is not None else None
        try:
            response = await self._client.post(
                "/api/v1/knowledge-search",
                headers=headers,
                json={
                    "query": query,
                    "knowledge_base_ids": [str(knowledge_base_id)],
                },
            )
        except (httpx.TimeoutException, httpx.NetworkError) as exc:
            raise WeKnoraUnavailableError("WeKnora unavailable") from exc
        if not response.is_success:
            raise WeKnoraBadResponseError("invalid WeKnora response")
        try:
            payload = response.json()
        except ValueError as exc:
            raise WeKnoraBadResponseError("invalid WeKnora response") from exc
        if not isinstance(payload, dict) or payload.get("success") is not True:
            raise WeKnoraBadResponseError("invalid WeKnora response")
        data = payload.get("data")
        if not isinstance(data, list) or any(not isinstance(item, dict) for item in data):
            raise WeKnoraBadResponseError("invalid WeKnora response")
        return data