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
2026-08-29 08:57:18 +00:00
parent 58488de7b3
commit 6e2e7386ad
7 changed files with 827 additions and 12 deletions

View File

@@ -0,0 +1,142 @@
"""Integration tests for Stage 20: Context Budgeting with ResearchOrchestrator."""
from unittest.mock import MagicMock, patch
import pytest
from nsct.orchestration.context_budget import (
ContextBudgetConfig,
ContextBudgetExhaustedError,
ContextBudgetTracker,
)
from nsct.orchestration.orchestrator import ResearchOrchestrator
from nsct.config import AppSettings, AudioConfig, DatabaseConfig, LLMConfig, VisionConfig
@pytest.fixture
def mock_config():
"""Minimal AppSettings for orchestrator tests."""
llm = LLMConfig(
base_url="http://test:8000/v1",
model="test-model",
max_concurrency=3,
)
vision = VisionConfig(
base_url="http://test:8001/v1",
model="test-vision",
)
audio = AudioConfig(
base_url="http://test:8002/v1",
model="test-audio",
)
postgres = DatabaseConfig(url="sqlite+aiosqlite:///")
return AppSettings(llm=llm, vision=vision, audio=audio, postgres=postgres)
def test_orchestrator_has_context_budget_tracker(mock_config):
"""Orchestrator muss context_budget_tracker als Property haben."""
orchestrator = ResearchOrchestrator(
config=mock_config,
research_id="test-research-id",
query="test query",
)
tracker = orchestrator.context_budget_tracker
assert isinstance(tracker, ContextBudgetTracker)
assert tracker.total_usage == 0
def test_orchestrator_custom_context_budget(mock_config):
"""Orchestrator kann ein custom ContextBudgetConfig akzeptieren."""
custom_config = ContextBudgetConfig(
planner_max_tokens=16000,
claim_extraction_max_tokens=24000,
contradiction_max_tokens=32000,
synthesis_max_tokens=64000,
)
orchestrator = ResearchOrchestrator(
config=mock_config,
research_id="test-research-id",
query="test query",
context_budget_config=custom_config,
)
assert orchestrator._context_budget_config.planner_max_tokens == 16000
assert orchestrator.context_budget_tracker._config.planner_max_tokens == 16000
def test_orchestrator_track_tokens(mock_config):
"""Orchestrator muss track_tokens für Stages unterstützen."""
orchestrator = ResearchOrchestrator(
config=mock_config,
research_id="test-research-id",
query="test query",
)
# Track planner tokens
orchestrator._track_context_tokens("planner", 5000)
assert orchestrator.context_budget_tracker._usage.get("planner", 0) == 5000
# Track again — accumulation works
orchestrator._track_context_tokens("planner", 3000)
assert orchestrator.context_budget_tracker._usage.get("planner", 0) == 8000
# get_usage returns correct data
usage = orchestrator.context_budget_tracker.get_usage()
assert usage["planner"]["usage_tokens"] == 8000.0
assert usage["planner"]["limit_tokens"] == 12000.0
def test_orchestrator_context_budget_exhaustion(mock_config):
"""Orchestrator wirft ContextBudgetExhaustedError bei Budget-Überschreitung."""
orchestrator = ResearchOrchestrator(
config=mock_config,
research_id="test-research-id",
query="test query",
)
# Track 12000 tokens to reach the limit
orchestrator._track_context_tokens("planner", 12000)
# 12000 tokens tracked → at limit → is_exhausted returns True
assert orchestrator.context_budget_tracker.is_exhausted("planner") is True
# One more token should raise
with pytest.raises(ContextBudgetExhaustedError) as exc_info:
orchestrator._track_context_tokens("planner", 1)
assert exc_info.value.stage_name == "planner"
assert exc_info.value.used_tokens == 12000
assert exc_info.value.limit_tokens == 12000
def test_orchestrator_reset_clears_context_budget(mock_config):
"""Orchestrator reset sollte auch context_budget_tracker zurücksetzen."""
orchestrator = ResearchOrchestrator(
config=mock_config,
research_id="test-research-id",
query="test query",
)
orchestrator._track_context_tokens("planner", 5000)
orchestrator._track_context_tokens("synthesis", 10000)
assert orchestrator.context_budget_tracker.total_usage == 15000
# Simulate reset of context budget tracker
orchestrator.context_budget_tracker.reset_all()
assert orchestrator.context_budget_tracker.total_usage == 0
def test_orchestrator_budget_usage_in_report(mock_config):
"""Context budget usage sollte im Report enthalten sein."""
orchestrator = ResearchOrchestrator(
config=mock_config,
research_id="test-research-id",
query="test query",
)
orchestrator._track_context_tokens("planner", 5000)
orchestrator._track_context_tokens("claim_extraction", 8000)
usage = orchestrator.context_budget_tracker.get_usage()
assert "planner" in usage
assert "claim_extraction" in usage
assert usage["planner"]["usage_tokens"] == 5000.0
assert usage["claim_extraction"]["usage_tokens"] == 8000.0
# synthesis and contradiction should be present but zero
assert usage["synthesis"]["usage_tokens"] == 0.0
assert usage["contradiction"]["usage_tokens"] == 0.0