orchestrata_ultimate.gateway.handlers.memory

Reference › orchestrata_ultimate › gateway › handlers › memory

MemoryHandler serves the memory.store, memory.retrieve, memory.search and memory.delete RPCs on top of the MemoryEngine (SPEC-010; see the memory package). rpc is its importer. The gateway injects its shared engine with set_engine, so the RPC surface and the session-close cold-migration path use one engine; without that, _ensure_engine builds a default engine and initialises it on first use. _require raises RPCError with an invalid-params code when a mandatory parameter is missing. handle_store needs content and session_id and saves to the hot store; handle_retrieve fetches by key, hot store first and then cold; handle_search runs a semantic search with limit and min_similarity parameters; handle_delete removes an entry from both stores. _entry_to_dict shapes entries for the reply. See gateway.

At a glance

Class diagram

Dependency diagram

API

Classes

MemoryHandler

Line 31

memory.* methods backed by the MemoryEngine.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/gateway/handlers/memory.py

"""Gateway handler exposing the memory engine (SPEC-010).

The ``MemoryEngine`` (Redis hot store + Neo4j cold store) is fully implemented but
was unreachable from the gateway — no handler, no ``memory.*`` RPC methods. This
handler registers them, delegating to a lazily-initialized engine (built and
connected on first use so it adds nothing to gateway startup). Tests inject a
test-scoped engine for isolation; production uses the env-configured default.
"""

from __future__ import annotations

import logging
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.gateway.exceptions import RPCError
from orchestrata_ultimate.memory.memory_engine import MemoryEngine

log = logging.getLogger(__name__)

_INVALID_PARAMS = -32602


@traces(spec="010", story="US-001", reqs=["REQ-010-004"])
def _entry_to_dict(entry: Any) -> dict[str, Any]:
    # Drop the embedding vector — large and not useful to RPC clients.
    dumped: dict[str, Any] = entry.model_dump(mode="json", exclude={"embedding"})
    return dumped


class MemoryHandler:
    """``memory.*`` methods backed by the ``MemoryEngine``."""

    def __init__(self, engine: MemoryEngine | None = None) -> None:
        """
        Args:
            engine: Optional pre-built engine (tests inject a scoped one). When
                omitted, a default env-configured engine is built lazily.
        """
        self._engine = engine
        self._initialized = bool(
            engine is not None and getattr(engine, "_running", False)
        )

    def set_engine(self, engine: MemoryEngine) -> None:
        """Inject the shared gateway-owned engine after construction.

        Used by the gateway to replace the lazily-built default with the one
        DI-registered singleton so the memory RPC surface and the session-close
        cold-migration path operate on the same engine (SPEC-010 FR-001).

        Args:
            engine: The gateway's singleton ``MemoryEngine``.
        """
        self._engine = engine
        self._initialized = getattr(engine, "_running", False)

    @traces(spec="010", story="US-001", reqs=["REQ-010-003"])
    async def _ensure_engine(self) -> MemoryEngine:
        if self._engine is None:
            self._engine = MemoryEngine()
        if not self._initialized:
            await self._engine.initialize({})
            self._initialized = True
        return self._engine

    @staticmethod
    def _require(params: dict[str, Any], name: str) -> Any:
        value = params.get(name)
        if value in (None, ""):
            raise RPCError(
                f"memory operation requires '{name}'", rpc_code=_INVALID_PARAMS
            )
        return value

    @traces(
        spec="010",
        story="US-001",
        reqs=["REQ-010-006", "REQ-010-009"],
    )
    async def handle_store(
        self, connection_id: str, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Handle memory.store — persist a memory entry to the hot store."""
        engine = await self._ensure_engine()
        entry = await engine.store(
            content=self._require(params, "content"),
            volk_id=params.get("volk_id", "HUMAN"),
            session_id=self._require(params, "session_id"),
            key=params.get("key"),
            tags=params.get("tags"),
            metadata=params.get("metadata"),
        )
        return {"entry": _entry_to_dict(entry)}

    @traces(spec="010", story="US-001", reqs=["REQ-010-011"])
    async def handle_retrieve(
        self, connection_id: str, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Handle memory.retrieve — fetch a single entry by key (hot then cold)."""
        engine = await self._ensure_engine()
        entry = await engine.retrieve(self._require(params, "key"))
        return {"entry": _entry_to_dict(entry) if entry else None}

    @traces(
        spec="010",
        story="US-004",
        reqs=["REQ-010-015", "REQ-010-019"],
    )
    async def handle_search(
        self, connection_id: str, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Handle memory.search — semantic search across stored memories."""
        engine = await self._ensure_engine()
        results = await engine.search(
            self._require(params, "query"),
            limit=int(params.get("limit", 10)),
            min_similarity=float(params.get("min_similarity", 0.7)),
        )
        return {"results": [_entry_to_dict(e) for e in results]}

    async def handle_delete(
        self, connection_id: str, params: dict[str, Any]
    ) -> dict[str, Any]:
        """Handle memory.delete — remove an entry from both stores."""
        engine = await self._ensure_engine()
        deleted = await engine.delete(self._require(params, "key"))
        return {"deleted": deleted}