stage20: context budgeting - per-stage token limits (planner 12k, claim 16k, contradiction 24k, synthesis 48k)
Implemented: - context_budget.py: ContextBudgetConfig (Pydantic, frozen) mit 4 Stage-Limits, validation (ge/le), get_limit(), total_max_tokens, stage_keys - context_budget.py: ContextBudgetTracker mit track_tokens(), get_usage(), is_exhausted(), reset_stage(), reset_all(), total_usage, elapsed_seconds - context_budget.py: ContextBudgetExhaustedError mit stage_name, used_tokens, limit_tokens - orchestrator.py: _track_context_tokens() Methode, context_budget_config/tracker init - orchestrator.py: context_budget_tracker property export - __init__.py: exports ContextBudgetConfig, ContextBudgetExhaustedError, ContextBudgetTracker - priority_queue.py: __aenter__/__aexit__ auf async geandert (Testfix) - test_context_budget.py: 28 Tests (DefaultConfig, TrackAccumulation, ExhaustedError, GetUsage, Reset, IsExhausted, InvalidStageNames, InvalidTokens) - test_context_budget_integration.py: 6 Tests (Orchestrator-Integration) - HANDOFF.md: Stage 20 abgeschlossen dokumentiert Tests: 34 passed (28 unit + 6 integration) + 14 performance = 48 total
This commit is contained in:
NSCT Agent
@@ -1,5 +1,15 @@
|
||||
"""NSCT Research Orchestrator (Stage 12)."""
|
||||
|
||||
from .context_budget import (
|
||||
ContextBudgetConfig,
|
||||
ContextBudgetExhaustedError,
|
||||
ContextBudgetTracker,
|
||||
)
|
||||
from .models import ResearchRun
|
||||
|
||||
__all__ = ["ResearchRun"]
|
||||
__all__ = [
|
||||
"ContextBudgetConfig",
|
||||
"ContextBudgetExhaustedError",
|
||||
"ContextBudgetTracker",
|
||||
"ResearchRun",
|
||||
]
|
||||
316
src/nsct/orchestration/context_budget.py
Normal file
316
src/nsct/orchestration/context_budget.py
Normal file
@@ -0,0 +1,316 @@
|
||||
"""Context-budget configuration and per-stage token tracking for Stage 20.
|
||||
|
||||
This module provides token-budget limits and runtime tracking that prevent
|
||||
individual pipeline stages from consuming more context than allotted.
|
||||
Each stage (e.g. planner, claim_extraction, contradiction, synthesis) has
|
||||
its own budget; exceeding the budget raises
|
||||
:exc:`ContextBudgetExhaustedError`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Custom exception
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ContextBudgetExhaustedError(Exception):
|
||||
"""Raised when a stage has exhausted its token budget."""
|
||||
|
||||
def __init__(self, stage_name: str, used: int, limit: int) -> None:
|
||||
self.stage_name = stage_name
|
||||
self.used_tokens = used
|
||||
self.limit_tokens = limit
|
||||
super().__init__(
|
||||
f"Context budget exhausted for stage {stage_name!r}: "
|
||||
f"used {used} tokens, limit {limit}."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ContextBudgetConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_K = int # shorthand for clarity in field defaults
|
||||
|
||||
|
||||
class ContextBudgetConfig(BaseModel, frozen=True):
|
||||
"""Immutable token-budget configuration for Stage 20 (Context Budgeting).
|
||||
|
||||
All fields are readonly via ``frozen=True``. Validation guarantees
|
||||
sensible minimum / maximum ranges so that a misconfigured budget
|
||||
cannot accidentally allow runaway context growth.
|
||||
"""
|
||||
|
||||
planner_max_tokens: int = Field(
|
||||
default=12_000,
|
||||
ge=8_000,
|
||||
le=16_000,
|
||||
description="Maximum tokens allocated to the planning stage.",
|
||||
)
|
||||
claim_extraction_max_tokens: int = Field(
|
||||
default=16_000,
|
||||
ge=8_000,
|
||||
le=24_000,
|
||||
description="Maximum tokens allocated to claim extraction.",
|
||||
)
|
||||
contradiction_max_tokens: int = Field(
|
||||
default=24_000,
|
||||
ge=16_000,
|
||||
le=32_000,
|
||||
description="Maximum tokens allocated to contradiction detection.",
|
||||
)
|
||||
synthesis_max_tokens: int = Field(
|
||||
default=48_000,
|
||||
ge=32_000,
|
||||
le=64_000,
|
||||
description="Maximum tokens allocated to synthesis.",
|
||||
)
|
||||
|
||||
@field_validator(
|
||||
"planner_max_tokens",
|
||||
"claim_extraction_max_tokens",
|
||||
"contradiction_max_tokens",
|
||||
"synthesis_max_tokens",
|
||||
mode="before",
|
||||
)
|
||||
@classmethod
|
||||
def _sanitize(cls, v: Any) -> int:
|
||||
"""Coerce int-like values (str, float) to ``int``."""
|
||||
if isinstance(v, str):
|
||||
return int(v)
|
||||
return int(v)
|
||||
|
||||
# -- convenience ---------------------------------------------------------
|
||||
|
||||
@property
|
||||
def stage_keys(self) -> list[str]:
|
||||
"""Return the list of stage-key names in field order."""
|
||||
return [
|
||||
"planner_max_tokens",
|
||||
"claim_extraction_max_tokens",
|
||||
"contradiction_max_tokens",
|
||||
"synthesis_max_tokens",
|
||||
]
|
||||
|
||||
@property
|
||||
def total_max_tokens(self) -> int:
|
||||
"""Sum of all per-stage budgets."""
|
||||
return (
|
||||
self.planner_max_tokens
|
||||
+ self.claim_extraction_max_tokens
|
||||
+ self.contradiction_max_tokens
|
||||
+ self.synthesis_max_tokens
|
||||
)
|
||||
|
||||
def get_limit(self, stage_key: str) -> int:
|
||||
"""Return the token limit for a given *stage_key* (e.g. ``'planner'``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
stage_key:
|
||||
One of the four stage identifiers: *planner*,
|
||||
*claim_extraction*, *contradiction*, or *synthesis*.
|
||||
|
||||
Returns
|
||||
-------
|
||||
int
|
||||
The token limit for the stage.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If the stage key is not recognised.
|
||||
"""
|
||||
field_name = f"{stage_key}_max_tokens"
|
||||
if not hasattr(self, field_name):
|
||||
raise ValueError(
|
||||
f"Unknown stage key {stage_key!r}; expected one of "
|
||||
f"{self.stage_keys}"
|
||||
)
|
||||
return getattr(self, field_name) # type: ignore[no-any-return]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ContextBudgetTracker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Maps the human-facing stage names used by ``track_tokens`` to config field
|
||||
# keys (the ``_max_tokens`` suffix is stripped).
|
||||
_STAGE_KEY_MAP: dict[str, str] = {
|
||||
"planner": "planner_max_tokens",
|
||||
"claim_extraction": "claim_extraction_max_tokens",
|
||||
"contradiction": "contradiction_max_tokens",
|
||||
"synthesis": "synthesis_max_tokens",
|
||||
}
|
||||
|
||||
# Approximate bytes-per-token ratio (used by ``get_usage`` for an estimate).
|
||||
_BYTES_PER_TOKEN: float = 4.0
|
||||
|
||||
|
||||
class ContextBudgetTracker:
|
||||
"""Tracks per-stage token consumption and enforces budget limits.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
config:
|
||||
A :class:`ContextBudgetConfig` instance defining the hard limits.
|
||||
"""
|
||||
|
||||
def __init__(self, config: ContextBudgetConfig) -> None:
|
||||
self._config = config
|
||||
self._start_time = time.monotonic()
|
||||
# stage_name -> tokens consumed (accumulated, never reset)
|
||||
self._usage: dict[str, int] = {}
|
||||
|
||||
# -- public API ----------------------------------------------------------
|
||||
|
||||
def track_tokens(self, stage_name: str, tokens: int) -> None:
|
||||
"""Record *tokens* consumed by *stage_name*.
|
||||
|
||||
If the stage's cumulative usage would exceed its configured limit
|
||||
a :exc:`ContextBudgetExhaustedError` is raised *before* the counter
|
||||
is incremented.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
stage_name:
|
||||
Human-readable stage name (e.g. ``'planner'``, ``'synthesis'``).
|
||||
tokens:
|
||||
Number of tokens consumed. Must be positive.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If *tokens* is not positive or *stage_name* is not recognised.
|
||||
ContextBudgetExhaustedError
|
||||
If the budget limit for this stage would be exceeded.
|
||||
"""
|
||||
if tokens <= 0:
|
||||
raise ValueError(f"tokens must be positive, got {tokens}")
|
||||
|
||||
field_key = _STAGE_KEY_MAP.get(stage_name)
|
||||
if field_key is None:
|
||||
known = ", ".join(_STAGE_KEY_MAP)
|
||||
raise ValueError(
|
||||
f"Unknown stage '{stage_name}'. Known stages: {known}"
|
||||
)
|
||||
|
||||
limit = getattr(self._config, field_key)
|
||||
current = self._usage.get(stage_name, 0)
|
||||
new_total = current + tokens
|
||||
|
||||
if new_total > limit:
|
||||
logger.error(
|
||||
"Context budget exceeded: stage=%s used=%d limit=%d (+%d)",
|
||||
stage_name, current, limit, tokens,
|
||||
)
|
||||
raise ContextBudgetExhaustedError(
|
||||
stage_name=stage_name, used=current, limit=limit
|
||||
)
|
||||
|
||||
self._usage[stage_name] = new_total
|
||||
logger.debug(
|
||||
"tracked %d tokens for stage %s (total %d / %d)",
|
||||
tokens, stage_name, new_total, limit,
|
||||
)
|
||||
|
||||
def get_usage(self) -> dict[str, dict[str, float]]:
|
||||
"""Return usage statistics for every stage.
|
||||
|
||||
Returns a flat dict keyed by stage name with the following keys:
|
||||
|
||||
- ``usage_tokens`` (int) — accumulated tokens consumed
|
||||
- ``limit_tokens`` (int) — configured maximum
|
||||
- ``usage_pct`` (float) — percentage of budget consumed (0-100)
|
||||
- ``usage_bytes`` (float) — estimated byte size (tokens × 4)
|
||||
|
||||
Example::
|
||||
|
||||
{
|
||||
"planner": {
|
||||
"usage_tokens": 9500,
|
||||
"limit_tokens": 12000,
|
||||
"usage_pct": 79.17,
|
||||
"usage_bytes": 38000.0,
|
||||
},
|
||||
...
|
||||
}
|
||||
"""
|
||||
result: dict[str, dict[str, float]] = {}
|
||||
for stage_name, field_key in _STAGE_KEY_MAP.items():
|
||||
limit = getattr(self._config, field_key)
|
||||
used = self._usage.get(stage_name, 0)
|
||||
pct = (used / limit * 100) if limit > 0 else 0.0
|
||||
result[stage_name] = {
|
||||
"usage_tokens": float(used),
|
||||
"limit_tokens": float(limit),
|
||||
"usage_pct": round(pct, 2),
|
||||
"usage_bytes": used * _BYTES_PER_TOKEN,
|
||||
}
|
||||
return result
|
||||
|
||||
def is_exhausted(self, stage_name: str | None = None) -> bool:
|
||||
"""Return ``True`` if any (optionally a specific) stage is exhausted.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
stage_name:
|
||||
If provided, only that stage is checked. Otherwise all stages
|
||||
are checked.
|
||||
"""
|
||||
stage_list = (
|
||||
[stage_name] if stage_name else list(_STAGE_KEY_MAP)
|
||||
)
|
||||
for s in stage_list:
|
||||
field_key = _STAGE_KEY_MAP[s]
|
||||
limit = getattr(self._config, field_key)
|
||||
used = self._usage.get(s, 0)
|
||||
if used >= limit:
|
||||
return True
|
||||
return False
|
||||
|
||||
def reset_stage(self, stage_name: str) -> None:
|
||||
"""Zero out the counter for *stage_name*.
|
||||
|
||||
Useful when restarting a stage after a failed run.
|
||||
"""
|
||||
if stage_name in self._usage:
|
||||
logger.info("Resetting token usage for stage %s", stage_name)
|
||||
self._usage[stage_name] = 0
|
||||
|
||||
def reset_all(self) -> None:
|
||||
"""Zero out all stage counters."""
|
||||
for key in self._usage:
|
||||
logger.info("Resetting token usage for stage %s", key)
|
||||
self._usage.clear()
|
||||
|
||||
@property
|
||||
def config(self) -> ContextBudgetConfig:
|
||||
"""Return the underlying :class:`ContextBudgetConfig`."""
|
||||
return self._config
|
||||
|
||||
@property
|
||||
def elapsed_seconds(self) -> float:
|
||||
"""Return wall-clock seconds since tracker creation."""
|
||||
return time.monotonic() - self._start_time
|
||||
|
||||
@property
|
||||
def total_usage(self) -> int:
|
||||
"""Sum of tokens consumed across all stages."""
|
||||
return sum(self._usage.values())
|
||||
|
||||
@property
|
||||
def total_limit(self) -> int:
|
||||
"""Sum of all configured stage limits."""
|
||||
return self._config.total_max_tokens
|
||||
@@ -30,6 +30,11 @@ from nsct.metrics import (
|
||||
)
|
||||
from nsct.models.claim import Claim
|
||||
from nsct.orchestration.budget import BudgetExhaustedError, BudgetTracker, HardBudgetConfig
|
||||
from nsct.orchestration.context_budget import (
|
||||
ContextBudgetConfig,
|
||||
ContextBudgetExhaustedError,
|
||||
ContextBudgetTracker,
|
||||
)
|
||||
from nsct.orchestration.models import ResearchRun
|
||||
from nsct.orchestration.state import ResearchRunState, StateMachine
|
||||
from nsct.providers.abstract import MultiProviderSearch, SearchProvider
|
||||
@@ -61,6 +66,7 @@ class ResearchOrchestrator:
|
||||
research_id: UUID,
|
||||
query: str,
|
||||
budget_config: HardBudgetConfig | None = None,
|
||||
context_budget_config: ContextBudgetConfig | None = None,
|
||||
depth: str = "normal",
|
||||
priority: Priority = Priority.NORMAL,
|
||||
) -> None:
|
||||
@@ -76,6 +82,8 @@ class ResearchOrchestrator:
|
||||
Die Forschungsfrage / Query.
|
||||
budget_config : HardBudgetConfig | None
|
||||
Optionales Budget. Wird aus config abgeleitet, wenn None.
|
||||
context_budget_config : ContextBudgetConfig | None
|
||||
Optionales Context-Budget. Wird aus config abgeleitet, wenn None.
|
||||
depth : str
|
||||
Suchtiefe ("quick", "normal", "deep").
|
||||
priority : Priority
|
||||
@@ -94,6 +102,13 @@ class ResearchOrchestrator:
|
||||
self._budget_config = HardBudgetConfig()
|
||||
self._budget_tracker = BudgetTracker(self._budget_config)
|
||||
|
||||
# Context Budget (Stage 20)
|
||||
if context_budget_config is not None:
|
||||
self._context_budget_config = context_budget_config
|
||||
else:
|
||||
self._context_budget_config = ContextBudgetConfig()
|
||||
self._context_budget_tracker = ContextBudgetTracker(self._context_budget_config)
|
||||
|
||||
# State Machine & Run
|
||||
self._state_machine = StateMachine(ResearchRunState.CREATED)
|
||||
self._run: ResearchRun | None = None
|
||||
@@ -131,6 +146,11 @@ class ResearchOrchestrator:
|
||||
"""BudgetTracker des aktuellen Runs."""
|
||||
return self._budget_tracker
|
||||
|
||||
@property
|
||||
def context_budget_tracker(self) -> ContextBudgetTracker:
|
||||
"""ContextBudgetTracker des aktuellen Runs (Stage 20)."""
|
||||
return self._context_budget_tracker
|
||||
|
||||
@property
|
||||
def is_completed(self) -> bool:
|
||||
"""True, wenn State == COMPLETED."""
|
||||
@@ -209,6 +229,24 @@ class ResearchOrchestrator:
|
||||
"""
|
||||
self._budget_tracker.check_budget()
|
||||
|
||||
def _track_context_tokens(self, stage_name: str, tokens: int) -> None:
|
||||
"""Tracke verbrauchte Token fuer Context Budgeting (Stage 20).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
stage_name : str
|
||||
Human-readable stage name (e.g. 'planner', 'claim_extraction',
|
||||
'contradiction', 'synthesis').
|
||||
tokens : int
|
||||
Number of tokens consumed.
|
||||
|
||||
Raises
|
||||
------
|
||||
ContextBudgetExhaustedError
|
||||
Wenn das Budget fuer die Stage erreicht ist.
|
||||
"""
|
||||
self._context_budget_tracker.track_tokens(stage_name, tokens)
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# Private: Sub-Component Init
|
||||
# ---------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user