tests/test_followup_rounds.py
Rama main · réplica de lectura del repositorio original.
"""
Integration test: a research run with several rounds actually searches again.
Round 1 searches the plan's terms. The coverage classifier reports missing
aspects and the continue decision asks for another round; round 2 must then
search for those aspects and fetch new sources.
"""
import json
import re
from tests.test_pipeline_smoke import _MockWebScraper, _ScriptedLLM # noqa: F401 (test isolation preamble)
from src.connectors.base import BaseConnector, ConnectorRegistry
from src.pipeline.models import SearchResult, SourceType
class _QueryAwareSearch(BaseConnector):
"""Returns one result per query, with a URL derived from the query."""
name = "query_aware_search"
def __init__(self):
self.queries: list[str] = []
async def search(self, query, max_results=10, **kwargs):
self.queries.append(query)
slug = re.sub(r"[^a-z0-9]+", "-", query.lower()).strip("-")
return [SearchResult(title=f"Result for {query}", url=f"https://example.org/{slug}",
snippet=f"About {query}.", source_type=SourceType.WEB_SEARCH)]
async def fetch(self, url):
return None
def can_handle(self, url):
return False
PLAN = {
"summary": "Power draw of the DGX B300",
"questions": [{
"id": "F1", "question": "How much power does the DGX B300 draw?",
"search_terms": {"en": ["DGX B300 power"]}, "search_langs": ["en"],
"source_scope": "web", "priority": "high",
}],
}
async def _run(max_rounds):
from src.config import PipelineConfig
from src.pipeline.orchestrator import ResearchOrchestrator
llm = _ScriptedLLM(
primary_responses={"You are a research planner": json.dumps(PLAN)},
harvest_responses={
"PERSON or ORGANISATION ANCHOR": json.dumps(
{"anchor_type": "none", "target": "", "confidence": 0.95, "reasoning": "product question"}),
"research question is answered by the": json.dumps(
{"coverage": "partial", "confidence": 0.8, "missing_aspects": ["idle power"],
"supporting_extract_ids": [], "reasoning": "peak power found, idle power missing"}),
"You control an autonomous research run": json.dumps(
{"decision": "continue", "confidence": 0.9, "next_round_focus": "idle power",
"reasoning": "idle power still missing"}),
"Diagnose the state of this research run": json.dumps(
{"diagnosis": "partial_success", "confidence": 0.8, "remediation": "",
"user_message": "ok", "reasoning": "-"}),
"Identify all FACTOIDS": "[]",
"extract relevant information from the following source": (
"[F1] [POSITIVE] Fact: The DGX B300 draws about 14 kW\n"
" Context: datasheet\n Reliability: high"),
},
)
search = _QueryAwareSearch()
registry = ConnectorRegistry()
registry.register(search, is_search_engine=True)
registry.register(_MockWebScraper({}), is_web_scraper=True)
config = PipelineConfig()
config.max_rounds = max_rounds
config.data_dir = "/tmp/followup_rounds_test"
orch = ResearchOrchestrator(llm=llm, connectors=registry, config=config)
async def progress(*_a, **_k):
pass
ctx = await orch.run(query="How much power does the DGX B300 draw?", chat_history=[],
context_docs="", template_name="", progress_callback=progress, mode="web")
return ctx, search
async def test_second_round_searches_for_the_missing_aspects():
ctx, search = await _run(max_rounds=2)
assert ctx.rounds_completed == 2
assert any("idle power" in q for q in search.queries), search.queries
urls = {s.url for s in ctx.sources}
assert any("idle-power" in u for u in urls), urls
async def test_single_round_configuration_does_not_search_again():
ctx, search = await _run(max_rounds=1)
assert ctx.rounds_completed == 1
assert not any("idle power" in q for q in search.queries)