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
143 lines
5.1 KiB
Python
143 lines
5.1 KiB
Python
"""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
|