opencode/NSCT---Neutral-Search-Crawler-Tool
Public Access
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

stage20: context budgeting - per-stage token limits (planner 12k, claim 16k, contradiction 24k, synthesis 48k)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

10
HANDOFF.md
View File

@@ -37,7 +37,7 @@ Webquellen recherchieren, Inhalte extrahieren, Quellen/Claims vergleichen, neutr
|||| **16** | ~~Observability~~ (✅ **ABGESCHLOSSEN** — `logging_config.py` + `metrics.py` + `test_logging.py` + `test_metrics.py`) |
||| **17** | ~~Neutralitäts-Tests~~ (✅ **ABGESCHLOSSEN** — `test_neutrality_a.py` + `test_neutrality_b.py` + `test_neutrality_c.py` + `test_neutrality_d.py` + `test_neutrality_e.py`) |
**Gesamt:** ~100 Dateien, ~17050 Zeilen Code, ~380 Tests.
**Gesamt:** ~105 Dateien, ~17400 Zeilen Code, ~428 Tests.
---
@@ -90,7 +90,7 @@ CREATED → PLANNING → SEARCHING → FETCHING → EXTRACTING → ANALYZING →
Wenn ein neuer Thread weiterarbeiten soll, einfach **Stage X** nennen und mit der Arbeit beginnen. Der neue Thread liest prompt.md (liegt im Repo als `/home/faligam/nsct/prompt.md`) für die volle Spezifikation und setzt bei der nächsten offenen Stage fort.
**Stage 19 ist die nächste offene Stage.**
**Stage 21 ist die nächste offene Stage.**
---
@@ -104,9 +104,9 @@ Wenn ein neuer Thread weiterarbeiten soll, einfach **Stage X** nennen und mit de
||| **15** | ~~CLI~~ (✅ **ABGESCHLOSSEN** — `cli.py` + `tests/test_cli.py`) |
|||| **16** | ~~Observability — Structured Logging, Metriken~~ (✅ **ABGESCHLOSSEN** — `logging_config.py` + `metrics.py` + `test_logging.py` + `test_metrics.py`) |
||| **17** | ~~Tests für Neutralitätsmethodik~~ (✅ **ABGESCHLOSSEN** — `test_neutrality_a_syndication.py` + `test_neutrality_b_political_statements.py` + `test_neutrality_c_scientific_disagreement.py` + `test_neutrality_d_prompt_injection.py` + `test_neutrality_e_missing_evidence.py`) |
||| **18** | ~~Docker Hardening~~ (✅ **ABGESCHLOSSEN** — `Dockerfile` hardened, `.dockerignore`, `security_opt`, health-check fix) |
|| **19** | ~~Performanceoptimierung — LLM Concurrency Semaphore(3), Priorisierung (HIGH/NORMAL/LOW), Batching~~ (✅ **ABGESCHLOSSEN** — `semaphore.py`, `priority_queue.py`, `llm.py`, `orchestrator.py` priority field, `test_performance.py`) |
| 20 | Context Budgeting — Pro Stage Kontext-Limits (Planner 8-16k, Claim 8-24k, Contradiction 16-32k, Synthesis 32-64k). |
|||| **18** | ~~Docker Hardening~~ (`6067908`) | 1 | – ||
||| **19** | ~~Performanceoptimierung — LLM Concurrency Semaphore(3), Priorisierung (HIGH/NORMAL/LOW)~~ (`4335a40`) | 5 | 14 ||
|| 20 | ~~Context Budgeting — Pro Stage Kontext-Limits (Planner 12k, Claim 16k, Contradiction 24k, Synthesis 48k)~~ (✅ `context_budget.py` + `orchestrator.py` context_budget + `_track_context_tokens()` + `__init__.py` + `test_context_budget.py` + `test_context_budget_integration.py` + `priority_queue.py` async_fix — 34 Tests) |
| 21 | Reproduzierbarkeit — research_run_hash, vollständige Provenance aller Schritte. |
| 22 | Abschluss & Production Readiness — README, ARCHITECTURE, SECURITY, METHODOLOGY, API, DEPLOYMENT. End-to-End-Test. |

12
src/nsct/orchestration/__init__.py
View File

@@ -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
View 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

38
src/nsct/orchestration/orchestrator.py
View File

@@ -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
# ---------------------------------------------------------------

19
tests/stages/test_performance.py
View File

@@ -85,10 +85,16 @@ class TestPriorityLimiter:
def test_release_only_for_non_high(self) -> None:
limiter = PriorityLimiter(max_concurrency=1)
limiter.acquire(Priority.HIGH)
async def test():
await limiter.acquire(Priority.HIGH)
# HIGH release is a no-op
limiter.release(Priority.HIGH)
# NORMAL still not acquired, so active_count should be 0
assert limiter.active_count() == 0
asyncio_run(test())
def test_max_concurrency_reflection(self) -> None:
limiter = PriorityLimiter(max_concurrency=5)
assert limiter.max_concurrency == 5
@@ -219,7 +225,7 @@ class TestLLMConcurrency:
async def test():
normal_done = False
async def slow():
async def slow(**kwargs):
nonlocal normal_done
await asyncio.sleep(0.1)
normal_done = True
@@ -311,5 +317,10 @@ class TestAsyncContextManager:
def test_context_manager(self) -> None:
limiter = PriorityLimiter(max_concurrency=1)
asyncio_run(limiter.__aenter__())
asyncio_run(limiter.__aexit__(None, None, None))
async def test():
ctx = await limiter.__aenter__()
assert ctx is limiter
await limiter.__aexit__(None, None, None)
asyncio_run(test())

298
tests/test_context_budget.py Normal file
View File

@@ -0,0 +1,298 @@
"""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)

142
tests/test_context_budget_integration.py Normal file
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