orchestrata_ultimate.memory.memory_engine

Reference › orchestrata_ultimate › memory › memory_engine

MemoryEngine is the single entry point to agent memory (SPEC-010). Volks and the pipeline never talk to Redis or the graph directly; the gateway (server and the memory handler), the kernel pipeline and the pipeline runner use this class. It implements the engine protocol from core.protocols, so it is also a lifecycle component with initialize, start, stop, get_health and get_metadata. The package story is in the memory narrative.

The constructor builds the collaborators: the Redis hot store (hot_store), the graph cold store (cold_store), an EmbeddingService, a word-based token counter and a ContextCompressor. Redis host and port come from the core Redis factory when not passed, the graph URI from the graph factory, and the graph carries no credentials. initialize connects both stores and is idempotent, because the gateway can reach it from more than one path and a second connect would leak a client. stop closes both stores; health is HEALTHY while running and STOPPED otherwise.

• store rejects content over the length limit or more tags than the limit with ValueError, computes the embedding and token count, and writes the entry to the hot store only.
• retrieve looks the key up across all sessions in the hot store, then falls back to the cold store.
• search embeds the query, pre-filters the hot store by scope, max_entries (default candidate cap 1000) and max_age_seconds, ranks candidates by cosine similarity above min_similarity, multiplies each score by a half-life decay that depends on the entry's classification, and returns the best limit entries by decayed score. It searches the hot layer only.
• summarize compresses a session's hot entries (cold entries if the hot layer holds none) into a token-bounded MemorySummary.
• update and delete change an entry; delete removes the key from both layers.
• migrate_to_cold copies every hot entry of a session to the graph and deletes it from Redis, logs failures per entry and returns the number migrated.

At a glance

Class diagram

Dependency diagram

API

Classes

MemoryEngine (protocol)

Bases: MemoryEngineProtocol · Decorators: traces(spec='010', story='US-001', reqs=['REQ-010-003']) · Line 48

Unified agent memory engine.

Source

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

"""Memory Engine — unified agent memory abstraction.

SPEC-010: Orchestrates hot storage (Redis), cold storage (graph),
embedding service, and context compression into a single interface.
"""

import logging
import math
import uuid as _uuid
from datetime import UTC, datetime, timedelta
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.graph_factory import resolve_kg_database
from orchestrata_ultimate.core.protocols import (
    ComponentMetadata,
    ComponentType,
    HealthStatus,
    MemoryClassification,
    MemoryEngineProtocol,
    MemoryEntry,
    MemoryScope,
    MemorySummary,
    SummarizationStrategy,
)
from orchestrata_ultimate.core.redis_factory import (
    resolve_redis_host,
    resolve_redis_port,
)
from orchestrata_ultimate.memory.cold_store import GraphMemoryStore
from orchestrata_ultimate.memory.compressor import ContextCompressor
from orchestrata_ultimate.memory.embedding import EmbeddingService
from orchestrata_ultimate.memory.hot_store import RedisMemoryStore
from orchestrata_ultimate.memory.token_counter import WordTokenCounter

logger = logging.getLogger(__name__)

MAX_CONTENT_LENGTH = 32 * 1024  # 32KB
MAX_TAG_COUNT = 50
DEFAULT_DECAY_HALF_LIFE = {
    MemoryClassification.PRIVATE: 24.0,
    MemoryClassification.SHARED: 168.0,
    MemoryClassification.SYSTEM: 720.0,
}


@traces(spec="010", story="US-001", reqs=["REQ-010-003"])
class MemoryEngine(MemoryEngineProtocol):
    """Unified agent memory engine.

    Orchestrates dual-layer storage (Redis hot / graph cold),
    embedding-based semantic search, and context compression.
    """

    @traces(spec="010", story="US-010-001", reqs=["REQ-010-002"])
    def __init__(
        self,
        redis_host: str | None = None,
        redis_port: int | None = None,
        redis_db: int = 0,
        redis_password: str | None = None,
        graph_uri: str | None = None,
        # Empty, not an env read: NEO4J_USERNAME/NEO4J_PASSWORD were retired
        # with the Neo4j containers on 2026-07-27 and every remaining graph
        # endpoint negotiates no auth. These were also DEFAULT ARGUMENTS, so
        # the env was sampled once at import and no later override could reach
        # them — a per-test or per-process change was silently ignored.
        graph_username: str = "",
        graph_password: str = "",
        graph_database: str | None = None,
    ):
        self._hot = RedisMemoryStore(
            redis_host=redis_host if redis_host is not None else resolve_redis_host(),
            redis_port=redis_port if redis_port is not None else resolve_redis_port(),
            redis_db=redis_db,
            redis_password=redis_password,
        )
        self._cold = GraphMemoryStore(
            uri=graph_uri,
            username=graph_username,
            password=graph_password,
            database=resolve_kg_database(graph_database),
        )
        self._embedding = EmbeddingService()
        self._counter = WordTokenCounter()
        self._compressor = ContextCompressor(self._counter)
        self._running = False
        self._metadata = ComponentMetadata(
            id="orchestrata_ultimate-memory",
            name="MemoryEngine",
            version="0.3.0",
            component_type=ComponentType.MEMORY_ENGINE,
            description="Unified agent memory with dual-layer storage",
            capabilities=["store", "retrieve", "search", "summarize", "migrate"],
        )

    @traces(spec="010", story="US-010-001", reqs=["REQ-010-003"])
    async def initialize(self, config: dict[str, Any]) -> None:
        """Initialize the memory engine with configuration.

        Idempotent: a second call while already running is a no-op. The
        underlying Redis/graph stores rebuild their client on every
        ``connect()`` (dropping the previous one without closing it), so an
        unguarded re-init would leak a connection. This matters now that a
        single gateway-owned engine can be initialized from more than one path
        — the DI resolve hook (which calls ``initialize`` on ``IComponent``
        instances) and the handler lazy-init.
        """
        if self._running:
            return
        await self._hot.connect()
        await self._cold.connect()
        self._running = True
        logger.info("MemoryEngine initialized")

    @traces(spec="010", story="US-010-001", reqs=["REQ-010-003"])
    async def start(self) -> None:
        """Start the memory engine."""
        self._running = True
        logger.info("MemoryEngine started")

    @traces(spec="010", story="US-010-001", reqs=["REQ-010-003"])
    async def stop(self) -> None:
        """Stop the memory engine and release resources."""
        self._running = False
        await self._hot.close()
        await self._cold.close()
        logger.info("MemoryEngine stopped")

    @traces(spec="010", story="US-010-001", reqs=["REQ-010-003"])
    async def get_health(self) -> HealthStatus:
        """Return health status."""
        if not self._running:
            return HealthStatus.STOPPED
        return HealthStatus.HEALTHY

    def get_metadata(self) -> ComponentMetadata:
        """Return component metadata."""
        return self._metadata

    @traces(
        spec="010",
        story="US-001",
        reqs=["REQ-010-006", "REQ-010-007", "REQ-010-008", "REQ-010-009"],
    )
    async def store(
        self,
        content: str,
        *,
        volk_id: str,
        session_id: str,
        key: str | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        classification: MemoryClassification = MemoryClassification.PRIVATE,
        task_id: str | None = None,
    ) -> MemoryEntry:
        """Store a memory entry."""
        if len(content) > MAX_CONTENT_LENGTH:
            raise ValueError(
                f"Content exceeds maximum length of {MAX_CONTENT_LENGTH} bytes"
            )
        tags = tags or []
        if len(tags) > MAX_TAG_COUNT:
            raise ValueError(f"Tag count exceeds maximum of {MAX_TAG_COUNT}")

        entry_key = key or str(_uuid.uuid4())
        now = datetime.now(UTC)
        embedding = self._embedding.compute_embedding(content)
        token_count = self._counter.count(content)
        entry = MemoryEntry(
            key=entry_key,
            content=content,
            tags=tags,
            metadata=metadata or {},
            classification=classification,
            embedding=embedding,
            created_at=now,
            updated_at=now,
            volk_id=volk_id,  # pyright: ignore[reportCallIssue]  # pydantic populate_by_name keyword not modeled by pyright
            session_id=session_id,
            task_id=task_id,
            token_count=token_count,
        )
        await self._hot.store(entry)
        logger.debug("Stored memory %s for %s", entry_key, volk_id)
        return entry

    @traces(spec="010", story="US-001", reqs=["REQ-010-002"])
    @traces(spec="017", story="US-001", reqs=["REQ-017-067"])
    async def retrieve(
        self,
        key: str,
        *,
        include_metadata: bool = True,
    ) -> MemoryEntry | None:
        """Retrieve a memory entry by exact key."""
        entry = await self._hot.retrieve_by_key_global(key)
        if entry:
            return entry
        entry = await self._cold.retrieve(key)
        return entry

    @traces(
        spec="010",
        story="US-004",
        reqs=["REQ-010-015", "REQ-010-017", "REQ-010-018", "REQ-010-019"],
    )
    async def search(
        self,
        query: str,
        *,
        scope: MemoryScope | None = None,
        limit: int = 10,
        min_similarity: float = 0.7,
        max_entries: int | None = None,
        max_age_seconds: int | None = None,
    ) -> list[MemoryEntry]:
        """Semantic search over memory entries.

        Pre-filters by ``max_entries`` and ``max_age_seconds`` before embedding
        similarity scoring so recall latency stays bounded as the hot store
        grows.
        """
        query_embedding = self._embedding.compute_embedding(query)
        scope = scope or MemoryScope()
        if max_age_seconds is not None:
            cutoff = datetime.now(UTC) - timedelta(seconds=max_age_seconds)
            if scope.date_from is None or cutoff > scope.date_from:
                scope.date_from = cutoff

        candidate_limit = max_entries if max_entries is not None else 1000
        hot_entries = await self._hot.retrieve_by_scope(
            scope,
            limit=candidate_limit,
            max_age_seconds=max_age_seconds,
        )
        candidates = []
        for entry in hot_entries:
            if entry.embedding:
                candidates.append((entry.key, entry.embedding))
            else:
                ekey_data = entry.key
                emb = self._embedding.compute_embedding(entry.content)
                candidates.append((ekey_data, emb))
        matches = self._embedding.search(
            query_embedding,
            candidates,
            limit=limit,
            min_similarity=min_similarity,
        )
        matched_keys = {k for k, _ in matches}
        results = []
        for entry in hot_entries:
            if entry.key in matched_keys:
                score = next(s for k, s in matches if k == entry.key)
                entry.relevance_score = score
                entry = self._apply_decay(entry)
                results.append(entry)
        results.sort(key=lambda e: e.relevance_score or 0.0, reverse=True)
        return results[:limit]

    @traces(spec="010", story="US-005", reqs=["REQ-010-020"])
    async def summarize(
        self,
        context_id: str,
        *,
        max_tokens: int = 2048,
        strategy: SummarizationStrategy = SummarizationStrategy.RELEVANCE,
    ) -> MemorySummary:
        """Compress memories for context_id into a token-bounded summary."""
        entries = await self._hot.get_all_entries(context_id)
        if not entries:
            cold_entries = await self._cold.search_by_scope(
                MemoryScope(session_id=context_id), limit=100
            )
            entries = cold_entries
        return self._compressor.summarize(
            entries, max_tokens=max_tokens, strategy=strategy
        )

    async def update(
        self,
        key: str,
        *,
        content: str | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        classification: MemoryClassification | None = None,
    ) -> MemoryEntry:
        """Update an existing memory entry."""
        entry = await self.retrieve(key)
        if entry is None:
            raise KeyError(f"Memory entry not found: {key}")
        if content is not None:
            if len(content) > MAX_CONTENT_LENGTH:
                raise ValueError(
                    f"Content exceeds maximum length of {MAX_CONTENT_LENGTH} bytes"
                )
            entry.content = content
            entry.embedding = self._embedding.compute_embedding(content)
            entry.token_count = self._counter.count(content)
        if tags is not None:
            if len(tags) > MAX_TAG_COUNT:
                raise ValueError(f"Tag count exceeds maximum of {MAX_TAG_COUNT}")
            entry.tags = tags
        if metadata is not None:
            entry.metadata = metadata
        if classification is not None:
            entry.classification = classification
        entry.updated_at = datetime.now(UTC)
        await self._hot.store(entry)
        return entry

    async def delete(self, key: str) -> bool:
        """Delete a memory entry by key."""
        entry = await self._hot.retrieve_by_key_global(key)
        if entry:
            result = await self._hot.delete(entry.session_id, entry.volk_id, key)
            await self._cold.delete(key)
            return result
        cold_result = await self._cold.delete(key)
        return cold_result

    @traces(spec="010", story="US-002", reqs=["REQ-010-032"])
    async def migrate_to_cold(self, session_id: str) -> int:
        """Migrate all hot memories for session to cold storage."""
        entries = await self._hot.get_all_entries(session_id)
        migrated = 0
        for entry in entries:
            try:
                await self._cold.store(entry)
                await self._hot.delete(entry.session_id, entry.volk_id, entry.key)
                migrated += 1
            except Exception as e:
                logger.error("Failed to migrate memory %s: %s", entry.key, e)
        logger.info(
            "Migrated %d/%d memories for session %s",
            migrated,
            len(entries),
            session_id,
        )
        return migrated

    def _apply_decay(self, entry: MemoryEntry) -> MemoryEntry:
        """Apply temporal decay to relevance score."""
        if entry.relevance_score is None:
            return entry
        half_life = DEFAULT_DECAY_HALF_LIFE.get(entry.classification, 24.0)
        now = datetime.now(UTC)
        age_hours = (now - entry.created_at).total_seconds() / 3600.0
        decay = math.exp(-math.log(2) * age_hours / half_life)
        entry.relevance_score = entry.relevance_score * decay
        return entry