Fix research pipeline budget tracking
This commit is contained in:
@@ -66,7 +66,7 @@ DEPTH_CONFIGS: dict[str, DepthConfig] = {
|
||||
max_pages_per_domain=2,
|
||||
max_total_download_bytes=500_000,
|
||||
max_llm_requests=15,
|
||||
max_research_duration_seconds=300,
|
||||
max_research_duration_seconds=1_800,
|
||||
max_context_per_llm_call=16_000,
|
||||
),
|
||||
"normal": DepthConfig(
|
||||
@@ -76,7 +76,7 @@ DEPTH_CONFIGS: dict[str, DepthConfig] = {
|
||||
max_pages_per_domain=5,
|
||||
max_total_download_bytes=2_000_000,
|
||||
max_llm_requests=30,
|
||||
max_research_duration_seconds=300,
|
||||
max_research_duration_seconds=3_600,
|
||||
max_context_per_llm_call=24_000,
|
||||
),
|
||||
"deep": DepthConfig(
|
||||
@@ -86,7 +86,7 @@ DEPTH_CONFIGS: dict[str, DepthConfig] = {
|
||||
max_pages_per_domain=10,
|
||||
max_total_download_bytes=5_000_000,
|
||||
max_llm_requests=60,
|
||||
max_research_duration_seconds=600,
|
||||
max_research_duration_seconds=7_200,
|
||||
max_context_per_llm_call=32_000,
|
||||
),
|
||||
}
|
||||
|
||||
@@ -167,7 +167,10 @@ class CrawlerManager:
|
||||
@staticmethod
|
||||
def _error_doc(url: str, error: str) -> NormalizedDocument:
|
||||
"""Create a NormalizedDocument representing an error."""
|
||||
doc = NormalizedDocument(
|
||||
# Keep the error result structurally identical to a successfully
|
||||
# normalized document. ``from_text`` hashes the (empty) content, so
|
||||
# downstream consumers can safely rely on content_hash being present.
|
||||
doc = NormalizedDocument.from_text(
|
||||
url=url,
|
||||
text="",
|
||||
title="",
|
||||
@@ -176,7 +179,6 @@ class CrawlerManager:
|
||||
"error": error,
|
||||
"content_type": "error",
|
||||
},
|
||||
content_hash="",
|
||||
extraction_tool="",
|
||||
)
|
||||
return doc
|
||||
|
||||
@@ -74,6 +74,7 @@ class BudgetTracker:
|
||||
def __init__(self, config: HardBudgetConfig) -> None:
|
||||
self._config = config
|
||||
self._start_time = time.monotonic()
|
||||
self._last_time_recorded = self._start_time
|
||||
# internal counters
|
||||
self._counters: dict[str, int | float] = {
|
||||
"search_queries": 0,
|
||||
@@ -132,12 +133,11 @@ class BudgetTracker:
|
||||
def record_time_elapsed(self) -> None:
|
||||
"""Tick the internal elapsed-time clock.
|
||||
|
||||
Adds the time since the tracker was created (or the last call)
|
||||
to the elapsed duration.
|
||||
Adds only the time since the previous call to the elapsed duration.
|
||||
"""
|
||||
self.update_research_duration(
|
||||
time.monotonic() - self._start_time
|
||||
)
|
||||
now = time.monotonic()
|
||||
self.update_research_duration(now - self._last_time_recorded)
|
||||
self._last_time_recorded = now
|
||||
|
||||
def increment_context_tokens(self, n: int) -> None:
|
||||
"""Add *n* tokens to the current LLM context window counter."""
|
||||
|
||||
@@ -345,8 +345,8 @@ class ResearchOrchestrator:
|
||||
ResearchRun
|
||||
Das erstellte Run-Objekt.
|
||||
"""
|
||||
# Budget init
|
||||
self._budget_tracker.increment_llm_requests(1) # planner call
|
||||
# The planner call itself is accounted for in ``_step_planning``.
|
||||
# Starting a run performs no LLM request.
|
||||
self._budget_tracker.increment_search(1) # initial search plan
|
||||
|
||||
# ResearchRun erstellen
|
||||
@@ -724,7 +724,26 @@ class ResearchOrchestrator:
|
||||
self._claims = []
|
||||
return {"success": True, "data": {"claims": []}, "claim_count": 0}
|
||||
|
||||
self._budget_tracker.increment_llm_requests(len(self._sources))
|
||||
# Error documents are retained in ``_sources`` for transparent
|
||||
# reporting, but they must never consume a claim-extraction LLM
|
||||
# request. Reserve one request for the synthesis step as well,
|
||||
# otherwise a quick run can use its entire LLM budget here and
|
||||
# fail deterministically before synthesis.
|
||||
extractable_sources = [source for source in self._sources if not source.get("error")]
|
||||
usage = self._budget_tracker.get_usage()
|
||||
remaining_llm_requests = self._budget_config.max_llm_requests - int(
|
||||
usage["max_llm_requests"]["usage"]
|
||||
)
|
||||
synthesis_reserve = 1
|
||||
max_sources_for_extraction = max(0, remaining_llm_requests - synthesis_reserve)
|
||||
extractable_sources = extractable_sources[:max_sources_for_extraction]
|
||||
|
||||
if not extractable_sources:
|
||||
logger.warning("No successful sources or LLM budget for claim extraction")
|
||||
self._claims = []
|
||||
return {"success": True, "data": {"claims": []}, "claim_count": 0}
|
||||
|
||||
self._budget_tracker.increment_llm_requests(len(extractable_sources))
|
||||
|
||||
# Stage 5: Claim Extraction
|
||||
try:
|
||||
@@ -734,7 +753,7 @@ class ResearchOrchestrator:
|
||||
llm_provider=llm_provider,
|
||||
config=self._config,
|
||||
research_run_id=self._run.id if self._run else uuid4(),
|
||||
sources=self._sources,
|
||||
sources=extractable_sources,
|
||||
)
|
||||
self._claims = await extractor.extract()
|
||||
except NameError:
|
||||
|
||||
Reference in New Issue
Block a user