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:
142
tests/test_context_budget_integration.py
Normal file
142
tests/test_context_budget_integration.py
Normal 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
|
||||
Reference in New Issue
Block a user