orchestrata_ultimate.context.budget

Reference › orchestrata_ultimate › context › budget

This module is the accounting backbone of the assembler. It does no retrieval. It serves spec 106 (REQ-106-030) and is tested by tests/context/test_budget.py.

ContextBudget is a frozen dataclass with a total window and one reservation for each layer: system, mission, plan, memory, artifacts, graph, ast, semantic, task, user and output. Construction raises ValueError when the reservations add up to more than the total. The properties reserved and available give the sum and the unreserved remainder, and layer_budget returns one layer's reservation, raising ValueError for an unknown layer name.

ContextUsage tracks actual spend against a budget. spend records tokens for a layer and raises ValueError for a negative amount, an unknown layer, a layer going over its own reservation, or the total window being exceeded. total_spent and remaining are derived properties. to_dict returns total, spent and remaining, plus a budget and spent pair for each layer, which is the shape written by usage_log.

At a glance

Class diagram

Dependency diagram

API

Classes

ContextBudget (dataclass)

Decorators: traces(spec='106', reqs=['REQ-106-030']), dataclass(frozen=True) · Line 18

Hard per-layer token budget for one LLM call.

ContextUsage (dataclass)

Decorators: dataclass · Line 75

Runtime spend tracker for a ContextBudget.

Source

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

"""Token-budget model for layered agent-context assembly.

A ``ContextBudget`` defines the maximum tokens each context layer may consume.
A ``ContextUsage`` tracks actual spend against that budget and refuses
over-allocation. This is the accounting backbone of the
``UnifiedContextAssembler``; it does not perform any retrieval itself.
"""

from __future__ import annotations

from dataclasses import dataclass, fields

from orchestrata_ultimate.core.base.traces import traces


@traces(spec="106", reqs=["REQ-106-030"])
@dataclass(frozen=True)
class ContextBudget:
    """Hard per-layer token budget for one LLM call.

    Attributes:
        total: Total model-context window available for this call.
        system: Tokens reserved for the volk system prompt.
        mission: Tokens reserved for session/goal/handoff state.
        plan: Tokens reserved for the current plan window.
        memory: Tokens reserved for memory summary + recall.
        artifacts: Tokens reserved for prior-stage artifact digest.
        graph: Tokens reserved for KG code-structure context.
        ast: Tokens reserved for AST-derived signals.
        semantic: Tokens reserved for semantic-search chunk context.
        task: Tokens reserved for task-type context pack.
        user: Tokens reserved for the user's raw instruction.
        output: Tokens reserved for the model's response.
    """

    total: int
    system: int = 0
    mission: int = 0
    plan: int = 0
    memory: int = 0
    artifacts: int = 0
    graph: int = 0
    ast: int = 0
    semantic: int = 0
    task: int = 0
    user: int = 0
    output: int = 0

    def __post_init__(self) -> None:
        reserved = sum(getattr(self, f.name) for f in fields(self) if f.name != "total")
        if reserved > self.total:
            raise ValueError(
                f"Context budget over-allocated: reserved {reserved} > "
                f"total {self.total}"
            )

    @property
    def reserved(self) -> int:
        """Sum of all reserved layer budgets."""
        return sum(getattr(self, f.name) for f in fields(self) if f.name != "total")

    @property
    def available(self) -> int:
        """Unreserved tokens that can be reallocated or held in reserve."""
        return self.total - self.reserved

    def layer_budget(self, layer: str) -> int:
        """Return the budget for a named layer."""
        if layer not in {f.name for f in fields(self)}:
            raise ValueError(f"Unknown context layer: {layer}")
        return int(getattr(self, layer))


@dataclass
class ContextUsage:
    """Runtime spend tracker for a ``ContextBudget``.

    ``spend`` raises ``ValueError`` if a layer would exceed its budget.
    The ``user`` layer is typically tracked last and may consume any
    remaining headroom not used by other layers.
    """

    budget: ContextBudget
    system: int = 0
    mission: int = 0
    plan: int = 0
    memory: int = 0
    artifacts: int = 0
    graph: int = 0
    ast: int = 0
    semantic: int = 0
    task: int = 0
    user: int = 0
    output: int = 0

    @property
    def total_spent(self) -> int:
        """Total tokens spent across all layers."""
        return sum(
            getattr(self, f.name)
            for f in fields(self)
            if f.name not in {"budget", "total"}
        )

    @property
    def remaining(self) -> int:
        """Tokens still available within the total budget."""
        return self.budget.total - self.total_spent

    def spend(self, layer: str, tokens: int) -> None:
        """Record ``tokens`` spent by ``layer``.

        Raises:
            ValueError: If ``layer`` is unknown, ``tokens`` is negative,
                or the spend would exceed the layer's budget or the total
                window.
        """
        if tokens < 0:
            raise ValueError(f"Cannot spend negative tokens: {tokens}")
        if layer not in {f.name for f in fields(self)} or layer == "budget":
            raise ValueError(f"Unknown context layer: {layer}")
        budget = self.budget.layer_budget(layer)
        current = int(getattr(self, layer))
        if current + tokens > budget:
            raise ValueError(
                f"Layer {layer} over budget: {current + tokens} > {budget}"
            )
        if self.total_spent + tokens > self.budget.total:
            raise ValueError(
                f"Total context window exceeded: {self.total_spent + tokens} > "
                f"{self.budget.total}"
            )
        object.__setattr__(self, layer, current + tokens)

    def to_dict(self) -> dict[str, int | dict[str, int]]:
        """Return spend and budget per layer for telemetry.

        Top-level keys (total/spent/remaining) are ``int``; each layer key
        (e.g. "graph", "user") is a ``{"budget": int, "spent": int}`` dict —
        the return type was previously (incorrectly) annotated as
        ``dict[str, int]``, which didn't match this nested shape.
        """
        layers = [f.name for f in fields(self) if f.name not in {"budget", "total"}]
        return {
            "total": self.budget.total,
            "spent": self.total_spent,
            "remaining": self.remaining,
            **{
                layer: {
                    "budget": self.budget.layer_budget(layer),
                    "spent": int(getattr(self, layer)),
                }
                for layer in layers
            },
        }