orchestrata_ultimate.context.assembler

Reference › orchestrata_ultimate › context › assembler

This module is the unified context assembler: it collects every source of prompt context for one volk call into a single bundle, trims each layer to a token reservation and reports what each layer spent. It does not talk to an LLM. It is used by orchestrata_ultimate.kernel.pipeline_kg and exported by the package __init__. Tests are tests/context/test_assembler.py, tests/kernel/test_pipeline_kg.py and tests/session/test_context_window.py. The class carries @traces for spec 106 (REQ-106-030, REQ-106-042). The package narrative is in context.

• ContextAssembly is a frozen result holding parts, a dictionary of text per layer, and usage, the ContextUsage from budget.
• UnifiedContextAssembler takes an optional ContextInjector, a ContextBudget, a token counter and a project id. Without a budget it uses _default_budget, a 128,000-token window with conservative per-layer reservations.

assemble builds the parts in a fixed order. The layers are system, mission, plan, memory, artifacts, graph, ast, semantic, task and user. The mission, plan, ast and semantic text come from the renderers in providers. Artifacts come from the store's own render_artifact_digest when it has one, with the generic provider as fallback if that fails or returns nothing. Memory is a thread summary plus a recall result. The graph and task layers come from the injector, when it accepts the subject. The injector receives a shallow copy of the payload because it mutates what it is given. Its output is cached per assembler instance under a key made of the graph snapshot, volk alias, subject, task type and the target-file fields of the payload. By default the snapshot id derives from the project id and the graph version, so a graph sync invalidates cached packs. The counters pack_cache_hits and pack_cache_misses expose cache behaviour.

Every layer except system and user is then trimmed with trim_to_budget, counted, and recorded with ContextUsage.spend. The usage dictionary is appended to a log through usage_log; a failed write is logged as a warning and does not stop assembly. Layers that used more than half their reservation are logged at info level. A spend that exceeds a layer budget raises ValueError.

At a glance

Class diagram

Dependency diagram

API

Classes

ContextAssembly (dataclass)

Decorators: dataclass(frozen=True) · Line 63

Result of assembling context for one volk call.

UnifiedContextAssembler

Decorators: traces(spec='106', reqs=['REQ-106-030', 'REQ-106-042']) · Line 71

Assemble layered context under a single token budget.

Functions

Source

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

"""Unified context assembler for volk prompts.

Collapses multiple context sources (KG code context, task-type packs, session
state, memory, artifacts, user intent) into one budgeted, token-accounted
bundle. The assembler does not talk to the LLM; it only prepares the prompt
material and reports how many tokens each layer consumed.

This is intentionally a skeleton: Phase 0 wires the existing KG injector and
user message; later phases add session, memory, artifact, AST, and plan-window
providers without changing the public ``assemble()`` contract.
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from orchestrata_ultimate.context.budget import (
    ContextBudget,
    ContextUsage,
)
from orchestrata_ultimate.context.providers.artifact_digest import (
    render_artifact_digest,
)
from orchestrata_ultimate.context.providers.ast_context import (
    render_ast_context,
)
from orchestrata_ultimate.context.providers.mission_context import (
    render_mission_context,
)
from orchestrata_ultimate.context.providers.plan_context import (
    render_plan_context,
)
from orchestrata_ultimate.context.providers.semantic_context import (
    render_semantic_context,
)
from orchestrata_ultimate.context.providers.trimming import (
    trim_to_budget,
)
from orchestrata_ultimate.context.usage_log import write_usage_record
from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.protocols import TokenCounterProtocol
from orchestrata_ultimate.knowledge_graph.context_injector import (
    ContextInjector,
    TaskType,
)
from orchestrata_ultimate.knowledge_graph.query_engine import get_graph_version
from orchestrata_ultimate.memory.memory_context import (
    recall_sync,
    summarize_for_thread_sync,
)
from orchestrata_ultimate.memory.token_counter import get_token_counter

if TYPE_CHECKING:
    from orchestrata_ultimate.session.models import ContextData

_logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class ContextAssembly:
    """Result of assembling context for one volk call."""

    parts: dict[str, str]
    usage: ContextUsage


@traces(spec="106", reqs=["REQ-106-030", "REQ-106-042"])
class UnifiedContextAssembler:
    """Assemble layered context under a single token budget.

    Args:
        injector: Optional ``ContextInjector`` for KG code/task context.
        budget: ``ContextBudget`` governing this assembler. Defaults to a
            128k-window budget with conservative per-layer reservations.
        counter: Token counter; defaults to the project-wide fallback.
    """

    def __init__(
        self,
        *,
        injector: ContextInjector | None = None,
        budget: ContextBudget | None = None,
        counter: TokenCounterProtocol | None = None,
        project_id: str | None = None,
    ) -> None:
        self._injector = injector
        self._budget = budget or _default_budget()
        self._counter = counter if counter is not None else get_token_counter()
        self._project_id = project_id
        # ContextPack cache: keyed by (snapshot, volk_alias, subject, task_type,
        # payload fingerprint). The cache lives for the lifetime of one assembler
        # instance (i.e. one pipeline run).
        self._pack_cache: dict[tuple[str, ...], dict[str, Any]] = {}
        self.pack_cache_hits = 0
        self.pack_cache_misses = 0

    def _default_snapshot_id(self) -> str:
        """Return a snapshot token tied to project_id and graph version."""
        project = self._project_id or ""
        return f"project={project}:gv={get_graph_version()}"

    def _pack_cache_key(
        self,
        *,
        payload: dict[str, Any],
        subject: str,
        task_type: TaskType | None,
        volk_alias: str,
        snapshot_id: str | None,
    ) -> tuple[str, ...]:
        """Build a stable cache key for a ContextPack render.

        The key includes the snapshot, volk, subject, task type, and a
        fingerprint of the payload fields that determine target files. It
        deliberately excludes the user message and session state, which vary
        per turn without changing the code-context pack.
        """
        files_created = tuple(sorted(payload.get("files_created", [])))
        project_dir = str(payload.get("project_dir", ""))
        file_tree = str(payload.get("file_tree", ""))
        return (
            snapshot_id if snapshot_id is not None else self._default_snapshot_id(),
            volk_alias,
            subject,
            str(task_type) if task_type is not None else "",
            *files_created,
            project_dir,
            file_tree,
        )

    def assemble(
        self,
        *,
        payload: dict[str, Any],
        subject: str,
        user_message: str,
        task_type: TaskType | None = None,
        task_params: dict[str, Any] | None = None,
        volk_alias: str = "",
        context_data: ContextData | None = None,
        artifact_store: Any | None = None,
        thread_id: str = "",
        snapshot_id: str | None = None,
    ) -> ContextAssembly:
        """Build context parts and token accounting for one volk call.

        The existing KG injector mutates the payload dict it receives, so a
        shallow copy is passed to keep the caller's payload intact.

        Args:
            payload: CXP packet payload.
            subject: CXP packet subject.
            user_message: The raw user/intent text for this volk.
            task_type: Optional task category for the task-type pack.
            task_params: Optional task-type parameters.
            volk_alias: Volk alias (cosmetic, passed to injector).
            context_data: Optional loaded session/project context for the
                mission block.
            artifact_store: Optional artifact store for prior-stage digest.
            thread_id: Pipeline thread id for artifact filtering.
            snapshot_id: Optional graph snapshot id. If omitted, a default
                snapshot derived from ``project_id`` and the current graph
                version is used so that post-sync calls invalidate cached packs.

        Returns:
            ``ContextAssembly`` with context parts and usage telemetry.
        """
        usage = ContextUsage(self._budget)
        parts: dict[str, str] = {
            "system": "",
            "mission": "",
            "plan": "",
            "memory": "",
            "artifacts": "",
            "graph": "",
            "ast": "",
            "semantic": "",
            "task": "",
            "user": user_message,
        }

        mission_text = render_mission_context(context_data)
        if mission_text:
            parts["mission"] = mission_text

        plan_text = render_plan_context(context_data)
        if plan_text:
            parts["plan"] = plan_text

        # Prefer the store's own ranked digest when available; otherwise fall back
        # to the generic provider that duck-types both kernel and orchestrator stores.
        artifact_text = ""
        if artifact_store is not None and hasattr(
            artifact_store, "render_artifact_digest"
        ):
            try:
                artifact_text = artifact_store.render_artifact_digest(
                    query=subject,
                    max_artifacts=10,
                    max_tokens=self._budget.layer_budget("artifacts"),
                )
            except Exception:
                _logger.exception("Store render_artifact_digest failed; falling back")
        if not artifact_text:
            artifact_text = render_artifact_digest(artifact_store, thread_id=thread_id)
        if artifact_text:
            parts["artifacts"] = artifact_text

        memory_parts: list[str] = []
        summary_text = summarize_for_thread_sync(thread_id, max_tokens=500)
        if summary_text:
            memory_parts.append(f"## Memory summary\n{summary_text}")
        # Reserve the remainder of the memory budget for recall; keep a floor so
        # we still surface at least a few tokens when the budget is tight.
        recall_budget = max(50, self._budget.layer_budget("memory") - 500)
        recall_text = recall_sync(
            thread_id,
            user_message,
            limit=3,
            max_tokens=recall_budget,
            max_entries=50,
            max_age_seconds=3600,
        )
        if recall_text:
            memory_parts.append(recall_text)
        if memory_parts:
            parts["memory"] = "\n\n".join(memory_parts)

        if self._injector is not None and self._injector.should_inject(subject):
            cache_key = self._pack_cache_key(
                payload=payload,
                subject=subject,
                task_type=task_type,
                volk_alias=volk_alias,
                snapshot_id=snapshot_id,
            )
            enriched = self._pack_cache.get(cache_key)
            if enriched is not None:
                self.pack_cache_hits += 1
            else:
                enriched = self._injector.inject(
                    payload.copy(),
                    task_type=task_type,
                    task_params=task_params,
                    volk_alias=volk_alias,
                )
                self._pack_cache[cache_key] = enriched
                self.pack_cache_misses += 1
            parts["graph"] = enriched.get("graph_context", "")
            task_context = enriched.get("task_context")
            if task_context is not None:
                parts["task"] = json.dumps(task_context, sort_keys=True)

        target_files = payload.get("files_created", [])
        ast_text = render_ast_context(
            target_files,
            subject=subject,
            project_id=self._project_id,
        )
        if ast_text:
            parts["ast"] = ast_text

        semantic_text = render_semantic_context(
            user_message=user_message,
            target_files=target_files,
            subject=subject,
            project_id=self._project_id,
        )
        if semantic_text:
            parts["semantic"] = semantic_text

        # Trim context layers at natural boundaries before accounting so a
        # single oversized layer does not blow up the whole assembly.
        # User/system/output are protected: user is the canonical instruction,
        # system is controlled by the volk, and output is reserved headroom.
        trimmable_layers = {
            "mission",
            "plan",
            "memory",
            "artifacts",
            "graph",
            "ast",
            "semantic",
            "task",
        }
        # "output" is reserved headroom, never an assembled part — the
        # trimmable_layers set alone decides what gets trimmed here.
        for layer in parts:
            text = parts[layer]
            if not text:
                continue
            if layer in trimmable_layers:
                budget = self._budget.layer_budget(layer)
                text = trim_to_budget(text, budget, self._counter)
                parts[layer] = text
            tokens = self._counter.count(text)
            usage.spend(layer, tokens)

        usage_dict = usage.to_dict()
        _logger.debug(
            "Context assembly usage: total=%d/%d, layers=%s",
            usage_dict["spent"],
            usage_dict["total"],
            usage_dict,
        )
        try:
            write_usage_record(
                usage_dict,
                volk_alias=volk_alias,
                subject=subject,
                thread_id=thread_id,
            )
        except Exception:
            _logger.warning("Failed to write context usage record", exc_info=True)
        # Surface any layer that consumed more than half its reservation.
        for layer in (
            "system",
            "mission",
            "plan",
            "memory",
            "artifacts",
            "graph",
            "ast",
            "semantic",
            "task",
            "user",
            "output",
        ):
            layer_info = usage_dict.get(layer)
            if not isinstance(layer_info, dict):
                continue
            spent = layer_info["spent"]
            budget = layer_info["budget"]
            if budget and spent > budget // 2:
                _logger.info("Context layer %r used %d/%d tokens", layer, spent, budget)

        return ContextAssembly(parts=parts, usage=usage)


def _default_budget() -> ContextBudget:
    """A conservative default budget for a 128k-context model."""
    return ContextBudget(
        total=128_000,
        system=800,
        mission=400,
        plan=600,
        memory=700,
        artifacts=500,
        graph=2_000,
        ast=500,
        semantic=500,
        task=1_000,
        user=4_000,
        output=2_000,
    )