orchestrata_ultimate.memory.memory_context

Reference › orchestrata_ultimate › memory › memory_context

A run-scoped registry that lets code which has no handle on the engine, such as volk call paths, the context assembler and the pipelines, recall memory for the current thread. The pipeline runner and kernel pipeline bind a MemoryEngine to a thread id with register_memory; the volk-side memory-context module and the context assembler read it back.

Binding is reference-counted. register_memory increments a count per thread id and the newest engine becomes the active provider; unregister_memory decrements it and only when it reaches zero removes the provider and the cached recall and summary results of that thread. get_memory returns the bound engine or None.

recall is best-effort by design: it returns an empty string when no engine is bound, the query is empty or the search raises, so memory never breaks a volk call. Otherwise it calls the engine's search with the limit and pre-filters, renders the entries with format_memories as a block that starts with "Relevant prior context from memory:" and one truncated bullet per entry (capped by max_tokens when set), and caches the block per argument combination. summarize_for_thread returns a cached MemorySummary per thread and token budget. recall_sync and summarize_for_thread_sync wrap them for synchronous callers through _run_async_sync, which runs the coroutine on a persistent background event loop thread (waiting up to 30 seconds) or, when a loop is already running in the caller's thread, with asyncio.run in a worker thread.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Constants

Functions

Source

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

"""Run-scoped memory-context provider registry.

The Memory Engine (SPEC-010) stores every conversation turn but nothing ever
*retrieved* it — volks received zero memory context. This module closes that gap
the same way tracing did: a memory engine is bound per run by ``thread_id`` (the
sync-router bridges volks onto worker threads, so a ContextVar would not survive;
volks look the provider up by their ``current_thread``). :func:`recall` is
best-effort — no provider registered (unit tests, or memory not wired) yields an
empty string, so nothing changes for callers without a memory engine.

This module lives in ``orchestrata_ultimate.memory`` (Layer 2) so both the agent
layer and the knowledge-graph/context layer can use it without cross-package
peer imports.
"""

from __future__ import annotations

import asyncio
import concurrent.futures
import logging
import threading
from typing import Any, cast

from orchestrata_ultimate.core.protocols import MemorySummary
from orchestrata_ultimate.memory.token_counter import get_token_counter

_log = logging.getLogger(__name__)

# thread_id -> (engine, refcount). Same-thread rounds can overlap (ack-only
# dispatch made this the normal case, not an edge): each register() call is a
# binder claiming the provider slot for the duration of its round, and only
# the LAST unregister() (refcount reaching zero) actually unbinds. Without
# this, whichever round finishes first would unbind while a concurrent
# round's volk is still mid-LLM, silently starving its recall() (F4a).
_providers: dict[str, tuple[Any, int]] = {}
# Per-run recall cache keyed by (thread_id, query, limit, max_tokens,
# max_entries, max_age_seconds). recall() runs once per LLM call; within a run
# the stored memories are stable (turns are persisted only after the pipeline
# completes), so caching a query's result is safe and avoids a store round-trip
# per call. Cleared for a thread on unregister_memory.
_recall_cache: dict[tuple[str | None, str, int, int, int | None, int | None], str] = {}
# Per-run summary cache keyed by (thread_id, max_tokens). Summarization is
# expensive (compresses all hot+cold memories for the thread) and the result is
# stable for the lifetime of a pipeline run, so cache it until the thread is
# unregistered.
_summary_cache: dict[tuple[str, int], MemorySummary] = {}
_lock = threading.Lock()

_MAX_ITEM_CHARS = 500

# Persistent background event loop for synchronous callers. Async Redis/graph
# clients are bound to the loop that created them; a fresh ``asyncio.run`` per
# call closes that loop and leaves clients unusable. A single daemon loop avoids
# this and keeps the sync bridge (kernel CLI, assembler) usable.
_bg_loop: asyncio.AbstractEventLoop | None = None
_bg_thread: threading.Thread | None = None


def _ensure_bg_loop() -> asyncio.AbstractEventLoop:
    """Return a running background event loop, creating it if necessary."""
    global _bg_loop, _bg_thread
    with _lock:
        if _bg_loop is not None and _bg_loop.is_running():
            return _bg_loop
        new_loop = asyncio.new_event_loop()
        _bg_loop = new_loop

        def _run() -> None:
            # Close over the local `new_loop` (known non-None), not the
            # module global -- pyright can't narrow a global's type across
            # a nested-function boundary (it may be reassigned by the time
            # this callback runs), even though it's the same object here.
            asyncio.set_event_loop(new_loop)
            new_loop.run_forever()

        _bg_thread = threading.Thread(target=_run, daemon=True)
        _bg_thread.start()
        return new_loop


def _run_async_sync(coro: Any) -> Any:
    """Run an async coroutine from a sync context on a persistent background loop.

    When an event loop is already running in the caller's thread, fall back to
    ``asyncio.run`` in a thread-pool subprocess (mirrors
    ``postman.sync_router._run_async_handler``). Otherwise schedule the coroutine
    on the shared background loop and wait for the result.
    """
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = None
    if loop is not None and loop.is_running():
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(asyncio.run, coro).result()

    bg = _ensure_bg_loop()
    future = asyncio.run_coroutine_threadsafe(coro, bg)
    return future.result(timeout=30)


def register_memory(thread_id: str, engine: Any) -> None:
    """Bind a memory engine as the recall source for ``thread_id``.

    Reference-counted: concurrent binders on the same ``thread_id`` (two
    same-thread rounds in flight at once) each increment the count and the
    provider stays bound until every binder has called
    :func:`unregister_memory`. The NEWEST engine passed always wins as the
    active provider (mirrors the pre-refcount last-writer-wins behavior for
    the single-binder case), but the slot is not released until the count
    returns to zero.
    """
    if not thread_id or engine is None:
        return
    with _lock:
        _, count = _providers.get(thread_id, (None, 0))
        _providers[thread_id] = (engine, count + 1)


def get_memory(thread_id: str | None) -> Any | None:
    """Return the memory engine bound to ``thread_id`` (or ``None``)."""
    if not thread_id:
        return None
    with _lock:
        entry = _providers.get(thread_id)
        return entry[0] if entry is not None else None


def unregister_memory(thread_id: str) -> None:
    """Release one binder's claim on ``thread_id`` (idempotent).

    Decrements the reference count and only pops the provider (and drops its
    caches) once the count reaches zero — i.e. once every concurrent binder
    for this thread has unregistered. Calling this more times than
    :func:`register_memory` was called is a no-op past zero.
    """
    with _lock:
        entry = _providers.get(thread_id)
        if entry is None:
            return
        _, count = entry
        if count > 1:
            engine, _ = entry
            _providers[thread_id] = (engine, count - 1)
            return
        del _providers[thread_id]
        for rkey in [k for k in _recall_cache if k[0] == thread_id]:
            del _recall_cache[rkey]
        for skey in [k for k in _summary_cache if k[0] == thread_id]:
            del _summary_cache[skey]


def format_memories(entries: list[Any], limit: int, max_tokens: int = 0) -> str:
    """Render memory entries into an injectable context block (or "").

    Args:
        entries: Memory records (objects or dicts with a ``content`` field).
        limit: Maximum number of entries to include.
        max_tokens: Optional token cap applied to the rendered block.
            ``0`` or negative means no cap.
    """
    lines: list[str] = []
    for entry in entries[:limit]:
        content = getattr(entry, "content", None)
        if content is None and isinstance(entry, dict):
            content = entry.get("content")
        text = str(content or "").strip()
        if text:
            lines.append(f"- {text[:_MAX_ITEM_CHARS]}")
    if not lines:
        return ""
    block = "Relevant prior context from memory:\n" + "\n".join(lines)
    if max_tokens > 0:
        counter = get_token_counter()
        if counter.count(block) > max_tokens:
            block = counter.truncate(block, max_tokens)
    return block


async def recall(
    thread_id: str | None,
    query: str,
    limit: int = 5,
    max_tokens: int = 0,
    max_entries: int | None = None,
    max_age_seconds: int | None = None,
) -> str:
    """Retrieve the memories most relevant to ``query`` as an injectable block.

    Best-effort: returns "" when no engine is registered for ``thread_id``, the
    query is empty, or retrieval fails — memory recall never breaks a volk call.

    Args:
        limit: Maximum results to return after similarity scoring.
        max_tokens: Optional token cap for the rendered block. ``0`` means no cap.
        max_entries: Pre-filter — maximum candidate entries to retrieve before
            similarity scoring.
        max_age_seconds: Pre-filter — only retrieve entries created within this
            many seconds of now.
    """
    engine = get_memory(thread_id)
    if engine is None or not query:
        return ""
    cache_key = (thread_id, query, limit, max_tokens, max_entries, max_age_seconds)
    with _lock:
        cached = _recall_cache.get(cache_key)
    if cached is not None:
        return cached
    try:
        entries = await engine.search(
            query,
            limit=limit,
            max_entries=max_entries,
            max_age_seconds=max_age_seconds,
        )
    except Exception as exc:  # retrieval is auxiliary — never break the volk
        _log.debug("memory recall failed for thread %s: %s", thread_id, exc)
        return ""
    block = format_memories(list(entries or []), limit, max_tokens=max_tokens)
    with _lock:
        _recall_cache[cache_key] = block
    return block


def recall_sync(
    thread_id: str | None,
    query: str,
    limit: int = 5,
    max_tokens: int = 0,
    max_entries: int | None = None,
    max_age_seconds: int | None = None,
) -> str:
    """Synchronous wrapper for :func:`recall`.

    Best-effort: returns "" when no engine is registered, the query is empty,
    or retrieval fails.

    Args:
        limit: Maximum results to return after similarity scoring.
        max_tokens: Optional token cap for the rendered block. ``0`` means no cap.
        max_entries: Pre-filter — maximum candidate entries to retrieve before
            similarity scoring.
        max_age_seconds: Pre-filter — only retrieve entries created within this
            many seconds of now.
    """
    if not thread_id or not query:
        return ""
    return str(
        _run_async_sync(
            recall(
                thread_id,
                query,
                limit,
                max_tokens,
                max_entries=max_entries,
                max_age_seconds=max_age_seconds,
            )
        )
    )


async def summarize_for_thread(
    thread_id: str | None, max_tokens: int = 512
) -> MemorySummary:
    """Return a token-bounded summary of all memories for ``thread_id``.

    The result is cached per ``(thread_id, max_tokens)`` until the thread is
    unregistered, because a pipeline run's memory set is stable and
    summarization touches every stored entry.
    """
    empty = MemorySummary(content="", token_count=0, source_count=0)
    if not thread_id:
        return empty
    cache_key = (thread_id, max_tokens)
    with _lock:
        cached = _summary_cache.get(cache_key)
    if cached is not None:
        return cached
    engine = get_memory(thread_id)
    if engine is None:
        return empty
    try:
        summary = cast(
            MemorySummary, await engine.summarize(thread_id, max_tokens=max_tokens)
        )
    except Exception as exc:
        _log.debug("memory summarize failed for thread %s: %s", thread_id, exc)
        return empty
    with _lock:
        _summary_cache[cache_key] = summary
    return summary


def summarize_for_thread_sync(thread_id: str | None, max_tokens: int = 512) -> str:
    """Synchronous wrapper returning the summarized text for ``thread_id``.

    Best-effort: returns "" when no engine is registered or summarization fails.
    """
    if not thread_id:
        return ""
    try:
        summary = _run_async_sync(summarize_for_thread(thread_id, max_tokens))
        return getattr(summary, "content", "") or ""
    except Exception as exc:
        _log.debug("memory summarize sync failed for thread %s: %s", thread_id, exc)
        return ""