6e2e7386ad
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
298 lines
10 KiB
Python
298 lines
10 KiB
Python
"""Tests for nsct.orchestration.context_budget – Stage 20."""
|
||
|
||
import pytest
|
||
|
||
from nsct.orchestration.context_budget import (
|
||
ContextBudgetConfig,
|
||
ContextBudgetExhaustedError,
|
||
ContextBudgetTracker,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper – build a tracker with given overrides; min-values from Field(ge/le)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_DEFAULTS = {
|
||
"planner_max_tokens": 12_000,
|
||
"claim_extraction_max_tokens": 16_000,
|
||
"contradiction_max_tokens": 24_000,
|
||
"synthesis_max_tokens": 48_000,
|
||
}
|
||
|
||
|
||
def _tracker(**overrides):
|
||
"""Return a fresh ContextBudgetTracker with optional budget overrides."""
|
||
merged = {**_DEFAULTS, **overrides}
|
||
return ContextBudgetTracker(ContextBudgetConfig(**merged))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1) Default-Konfiguration validiert sich selbst
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestDefaultConfig:
|
||
"""ContextBudgetConfig self-validation."""
|
||
|
||
def test_defaults_sane(self):
|
||
cfg = ContextBudgetConfig()
|
||
# Defaults lie inside their [ge, le] windows
|
||
assert cfg.planner_max_tokens == 12_000
|
||
assert cfg.claim_extraction_max_tokens == 16_000
|
||
assert cfg.contradiction_max_tokens == 24_000
|
||
assert cfg.synthesis_max_tokens == 48_000
|
||
# total is the sum
|
||
assert cfg.total_max_tokens == 12_000 + 16_000 + 24_000 + 48_000
|
||
|
||
def test_frozen(self):
|
||
"""Config is immutable – direct assignment is blocked."""
|
||
cfg = ContextBudgetConfig()
|
||
# Pydantic v2 frozen allows model_copy but not __setattr__
|
||
with pytest.raises(Exception):
|
||
cfg.planner_max_tokens = 999
|
||
|
||
def test_total_max_tokens(self):
|
||
cfg = ContextBudgetConfig()
|
||
assert cfg.total_max_tokens == 100_000
|
||
|
||
def test_get_limit_valid_key(self):
|
||
cfg = ContextBudgetConfig()
|
||
assert cfg.get_limit("planner") == 12_000
|
||
assert cfg.get_limit("synthesis") == 48_000
|
||
|
||
def test_get_limit_invalid_key_raises_value_error(self):
|
||
cfg = ContextBudgetConfig()
|
||
with pytest.raises(ValueError, match="Unknown stage key"):
|
||
cfg.get_limit("nonexistent")
|
||
|
||
def test_stage_keys(self):
|
||
cfg = ContextBudgetConfig()
|
||
keys = cfg.stage_keys
|
||
assert keys == [
|
||
"planner_max_tokens",
|
||
"claim_extraction_max_tokens",
|
||
"contradiction_max_tokens",
|
||
"synthesis_max_tokens",
|
||
]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2) track_tokens akkumuliert korrekt über mehrere calls hinweg
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestTrackAccumulation:
|
||
"""Accumulation of tracked tokens across calls."""
|
||
|
||
def test_accumulates(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
tr.track_tokens("planner", 200)
|
||
assert tr._usage["planner"] == 300
|
||
|
||
def test_multiple_stages_independent(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 50)
|
||
tr.track_tokens("claim_extraction", 30)
|
||
assert tr._usage["planner"] == 50
|
||
assert tr._usage["claim_extraction"] == 30
|
||
assert tr._usage["claim_extraction"] != tr._usage["planner"]
|
||
|
||
def test_total_usage_property(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
tr.track_tokens("synthesis", 200)
|
||
assert tr.total_usage == 300
|
||
|
||
def test_usage_starts_zero(self):
|
||
tr = _tracker()
|
||
assert tr._usage.get("planner", 0) == 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3) ContextBudgetExhaustedError wird bei Überschreitung geworfen
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestExhaustedError:
|
||
"""Raising ContextBudgetExhaustedError on budget overrun."""
|
||
|
||
def test_exact_limit_succeeds(self):
|
||
"""Hitting the limit exactly should still be allowed."""
|
||
tr = _tracker(planner_max_tokens=8_000) # minimum allowed
|
||
tr.track_tokens("planner", 8_000)
|
||
assert tr._usage["planner"] == 8_000
|
||
|
||
def test_over_limit_raises(self):
|
||
tr = _tracker(planner_max_tokens=8_000)
|
||
tr.track_tokens("planner", 4_000)
|
||
with pytest.raises(ContextBudgetExhaustedError) as exc_info:
|
||
tr.track_tokens("planner", 4_001) # 4000 + 4001 = 8001 > 8000
|
||
assert "planner" in str(exc_info.value)
|
||
assert exc_info.value.used_tokens == 4_000
|
||
assert exc_info.value.limit_tokens == 8_000
|
||
|
||
def test_over_limit_error_not_accumulated(self):
|
||
"""The counter must NOT be incremented on an overrun."""
|
||
tr = _tracker(planner_max_tokens=8_000)
|
||
tr.track_tokens("planner", 3_000)
|
||
with pytest.raises(ContextBudgetExhaustedError):
|
||
tr.track_tokens("planner", 5_001)
|
||
assert tr._usage["planner"] == 3_000 # unchanged
|
||
|
||
def test_error_attributes(self):
|
||
tr = _tracker(planner_max_tokens=8_000)
|
||
with pytest.raises(ContextBudgetExhaustedError) as exc_info:
|
||
tr.track_tokens("planner", 8_001)
|
||
err = exc_info.value
|
||
assert err.stage_name == "planner"
|
||
assert err.used_tokens == 0
|
||
assert err.limit_tokens == 8_000
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4) get_usage() liefert korrekte Pct und Bytes
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestGetUsage:
|
||
"""get_usage() returns correct pct and bytes."""
|
||
|
||
def test_tracks_tokens_and_bytes(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
usage = tr.get_usage()
|
||
p = usage["planner"]
|
||
assert p["usage_tokens"] == 100.0
|
||
assert p["limit_tokens"] == 12_000.0
|
||
# 100 / 12000 * 100 = 0.8333… %
|
||
assert p["usage_pct"] == 0.83
|
||
# 100 * 4.0 = 400.0 bytes
|
||
assert p["usage_bytes"] == 400.0
|
||
|
||
def test_unknown_stage_shows_zero(self):
|
||
"""Stages that have never been tracked still appear with zeros."""
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
usage = tr.get_usage()
|
||
assert usage["synthesis"]["usage_tokens"] == 0.0
|
||
assert usage["synthesis"]["usage_pct"] == 0.0
|
||
assert usage["synthesis"]["usage_bytes"] == 0.0
|
||
|
||
def test_usage_pct_rounded(self):
|
||
tr = _tracker(planner_max_tokens=8_000)
|
||
tr.track_tokens("planner", 1)
|
||
usage = tr.get_usage()
|
||
# 1/8000 * 100 = 0.0125 → rounded to 0.01
|
||
assert usage["planner"]["usage_pct"] == 0.01
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5) reset_stage und reset_all funktionieren
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestReset:
|
||
"""reset_stage and reset_all work correctly."""
|
||
|
||
def test_reset_stage(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
tr.reset_stage("planner")
|
||
assert tr._usage["planner"] == 0
|
||
|
||
def test_reset_stage_unknown_ignored(self):
|
||
"""Resetting an unknown stage is a no-op, not an error."""
|
||
tr = _tracker()
|
||
tr.reset_stage("nonexistent")
|
||
assert tr._usage == {}
|
||
|
||
def test_reset_all(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
tr.track_tokens("synthesis", 200)
|
||
tr.reset_all()
|
||
assert tr._usage == {}
|
||
assert tr.total_usage == 0
|
||
|
||
def test_after_reset_usage_reflects_zero(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
tr.reset_all()
|
||
usage = tr.get_usage()
|
||
assert usage["planner"]["usage_tokens"] == 0.0
|
||
assert usage["planner"]["usage_pct"] == 0.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 6) is_exhausted() erkennt exhaustion richtig
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestIsExhausted:
|
||
"""is_exhausted() detects exhaustion correctly."""
|
||
|
||
def test_not_exhausted_under_budget(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 100)
|
||
assert tr.is_exhausted("planner") is False
|
||
|
||
def test_exhausted_at_limit(self):
|
||
tr = _tracker(planner_max_tokens=8_000)
|
||
tr.track_tokens("planner", 8_000)
|
||
assert tr.is_exhausted("planner") is True
|
||
|
||
def test_is_exhausted_any_stage(self):
|
||
"""No argument → checks all stages."""
|
||
tr = _tracker()
|
||
tr.track_tokens("synthesis", 48_000) # at limit
|
||
assert tr.is_exhausted() is True
|
||
|
||
def test_is_exhausted_specific_false(self):
|
||
tr = _tracker(planner_max_tokens=8_000, synthesis_max_tokens=32_000)
|
||
tr.track_tokens("synthesis", 1_000)
|
||
# synthesis is under limit
|
||
assert tr.is_exhausted("synthesis") is False
|
||
|
||
def test_all_stages_clean(self):
|
||
tr = _tracker()
|
||
tr.track_tokens("planner", 1)
|
||
assert tr.is_exhausted() is False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 7) Ungültige Stage-Namen werfen ValueError
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestInvalidStageNames:
|
||
"""Invalid stage names raise ValueError."""
|
||
|
||
def track_tokens_bad_stage(self):
|
||
tr = _tracker()
|
||
with pytest.raises(ValueError, match="Unknown stage"):
|
||
tr.track_tokens("nonexistent", 100)
|
||
|
||
def get_limit_bad_stage(self):
|
||
cfg = ContextBudgetConfig()
|
||
with pytest.raises(ValueError, match="Unknown stage key"):
|
||
cfg.get_limit("nonexistent")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 8) Token = 0 oder negativ wirft ValueError
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestInvalidTokens:
|
||
"""Zero or negative token counts raise ValueError."""
|
||
|
||
def test_zero_tokens_raises(self):
|
||
tr = _tracker()
|
||
with pytest.raises(ValueError, match="tokens must be positive"):
|
||
tr.track_tokens("planner", 0)
|
||
|
||
def test_negative_tokens_raises(self):
|
||
tr = _tracker()
|
||
with pytest.raises(ValueError, match="tokens must be positive"):
|
||
tr.track_tokens("planner", -42) |