TencentDB-Agent-Memory/hermes-plugin/memory/memory_tencentdb/client.py

"""MemoryTencentdbSdkClient — HTTP client for the memory-tencentdb Gateway.

Wraps all Gateway API endpoints with timeout, retry, and error handling.
Thread-safe — can be shared across prefetch/sync threads.
"""

from __future__ import annotations

import json
import logging
import urllib.request
import urllib.error
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)

DEFAULT_TIMEOUT = 10  # seconds


class MemoryTencentdbSdkClient:
    """HTTP client for the memory-tencentdb Gateway sidecar."""

    def __init__(
        self,
        base_url: str = "http://127.0.0.1:8420",
        timeout: int = DEFAULT_TIMEOUT,
        api_key: Optional[str] = None,
    ):
        """Construct the client.

        Args:
            base_url: Gateway base URL.
            timeout: Default request timeout in seconds.
            api_key: Optional Bearer token. When non-empty, every request
                attaches ``Authorization: Bearer <api_key>``. When ``None``
                or empty, no auth header is sent — this preserves the
                pre-existing open-Gateway behaviour and is the right default
                for any deployment where the Gateway has not opted into
                ``TDAI_GATEWAY_API_KEY`` yet.

                The provider sources this value from
                ``MEMORY_TENCENTDB_GATEWAY_API_KEY`` (with
                ``TDAI_GATEWAY_API_KEY`` as a fallback). The Gateway must
                be configured with the matching secret independently —
                this client does not (and should not) propagate the value
                across to the Gateway process.
        """
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        # Strip whitespace defensively — env vars often pick up trailing
        # newlines from `echo` or YAML quoting; an exact-match Bearer
        # comparison would otherwise reject a key that "looks right".
        self._api_key = (api_key or "").strip() or None

    def _build_headers(self, *, content_type: bool) -> Dict[str, str]:
        """Build request headers, conditionally adding Authorization.

        Centralised so the auth header logic is stated once: every method
        below goes through ``_post`` / ``_get`` which call this helper. If
        you ever add a new HTTP verb, route it here.
        """
        headers: Dict[str, str] = {}
        if content_type:
            headers["Content-Type"] = "application/json"
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        return headers

    def _post(self, path: str, body: Dict[str, Any], timeout: Optional[int] = None) -> Dict[str, Any]:
        """Make a POST request to the Gateway."""
        url = f"{self._base_url}{path}"
        data = json.dumps(body).encode("utf-8")
        req = urllib.request.Request(
            url,
            data=data,
            headers=self._build_headers(content_type=True),
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=timeout or self._timeout) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            body_text = ""
            try:
                body_text = e.read().decode("utf-8", errors="replace")
            except Exception:
                pass
            logger.warning("memory-tencentdb Gateway %s returned %d: %s", path, e.code, body_text[:500])
            raise
        except Exception as e:
            logger.debug("memory-tencentdb Gateway %s failed: %s", path, e)
            raise

    def _get(self, path: str, timeout: Optional[int] = None) -> Dict[str, Any]:
        """Make a GET request to the Gateway."""
        url = f"{self._base_url}{path}"
        req = urllib.request.Request(
            url,
            headers=self._build_headers(content_type=False),
            method="GET",
        )
        try:
            with urllib.request.urlopen(req, timeout=timeout or self._timeout) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except Exception as e:
            logger.debug("memory-tencentdb Gateway GET %s failed: %s", path, e)
            raise

    # -- API methods ----------------------------------------------------------

    def health(self, timeout: int = 3) -> Dict[str, Any]:
        """Check if the Gateway is healthy."""
        return self._get("/health", timeout=timeout)

    def recall(self, query: str, session_key: str, user_id: str = "") -> Dict[str, Any]:
        """Recall memories for a query (prefetch)."""
        body: Dict[str, Any] = {"query": query, "session_key": session_key}
        if user_id:
            body["user_id"] = user_id
        return self._post("/recall", body)

    def capture(
        self,
        user_content: str,
        assistant_content: str,
        session_key: str,
        session_id: str = "",
        user_id: str = "",
    ) -> Dict[str, Any]:
        """Capture a conversation turn (sync_turn)."""
        body: Dict[str, Any] = {
            "user_content": user_content,
            "assistant_content": assistant_content,
            "session_key": session_key,
        }
        if session_id:
            body["session_id"] = session_id
        if user_id:
            body["user_id"] = user_id
        return self._post("/capture", body)

    def search_memories(self, query: str, limit: int = 5, type_filter: str = "", scene: str = "") -> Dict[str, Any]:
        """Search L1 structured memories."""
        body: Dict[str, Any] = {"query": query, "limit": limit}
        if type_filter:
            body["type"] = type_filter
        if scene:
            body["scene"] = scene
        return self._post("/search/memories", body)

    def search_conversations(self, query: str, limit: int = 5, session_key: str = "") -> Dict[str, Any]:
        """Search L0 raw conversations."""
        body: Dict[str, Any] = {"query": query, "limit": limit}
        if session_key:
            body["session_key"] = session_key
        return self._post("/search/conversations", body)

    def end_session(self, session_key: str, user_id: str = "") -> Dict[str, Any]:
        """End a session and trigger flush."""
        body: Dict[str, Any] = {"session_key": session_key}
        if user_id:
            body["user_id"] = user_id
        return self._post("/session/end", body)

    def seed(
        self,
        data: Any,
        session_key: str = "",
        strict_round_role: bool = False,
        auto_fill_timestamps: bool = True,
        config_override: Optional[Dict[str, Any]] = None,
        timeout: int = 300,
    ) -> Dict[str, Any]:
        """Batch seed historical conversations into the memory pipeline.

        Args:
            data: Seed input — Format A ``{"sessions": [...]}`` or Format B ``[...]``.
            session_key: Fallback session key when input sessions lack one.
            strict_round_role: Require each round to have both user and assistant.
            auto_fill_timestamps: Auto-fill missing timestamps (default True).
            config_override: Plugin config overrides (deep-merged).
            timeout: Request timeout in seconds (seed can be slow, default 300s).

        Returns:
            Summary dict with sessions_processed, rounds_processed, etc.
        """
        body: Dict[str, Any] = {"data": data}
        if session_key:
            body["session_key"] = session_key
        if strict_round_role:
            body["strict_round_role"] = True
        if not auto_fill_timestamps:
            body["auto_fill_timestamps"] = False
        if config_override:
            body["config_override"] = config_override
        return self._post("/seed", body, timeout=timeout)