251 lines
9.6 KiB
Python
251 lines
9.6 KiB
Python
"""Regression tests for the deployed backend research pipeline fixes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
from nsct.config import AppSettings, AudioConfig, DatabaseConfig, LLMConfig, VisionConfig
|
|
from nsct.orchestration.orchestrator import ResearchOrchestrator
|
|
from nsct.orchestration.budget import BudgetExhaustedError, BudgetTracker, HardBudgetConfig
|
|
from nsct.orchestration.state import ResearchRunState, StateMachine
|
|
from nsct.providers.abstract import NormalizedResult
|
|
from nsct.providers.llm import _LLMProviderImpl
|
|
from nsct.providers.metrics import ProviderMetrics
|
|
from nsct.providers.searxng import SearXNGProvider
|
|
from nsct.crawler.normalize import NormalizedDocument
|
|
from nsct.api.rest_research import DEPTH_CONFIGS
|
|
|
|
|
|
def _config(*, searxng_base_url: str | None = "http://searxng:8080") -> AppSettings:
|
|
return AppSettings(
|
|
llm=LLMConfig(base_url="http://llm.example/v1", model="test-model"),
|
|
vision=VisionConfig(base_url="http://vision.example/v1", model="vision"),
|
|
audio=AudioConfig(base_url="http://audio.example/v1"),
|
|
postgres=DatabaseConfig(url="sqlite+aiosqlite://"),
|
|
searxng_base_url=searxng_base_url,
|
|
)
|
|
|
|
|
|
def test_llm_base_url_does_not_duplicate_existing_v1() -> None:
|
|
provider = _LLMProviderImpl(_config(), ProviderMetrics())
|
|
provider._api_key = "test-key"
|
|
client = provider._create_client()
|
|
try:
|
|
assert str(client.base_url) == "http://llm.example/v1/"
|
|
finally:
|
|
asyncio.run(client.close())
|
|
|
|
|
|
def test_searxng_provider_normalizes_json_results() -> None:
|
|
async def run() -> None:
|
|
provider = SearXNGProvider("http://searxng:8080")
|
|
provider._client = httpx.AsyncClient(
|
|
transport=httpx.MockTransport(
|
|
lambda request: httpx.Response(
|
|
200,
|
|
json={
|
|
"results": [
|
|
{"url": "https://example.org/a", "title": "A", "content": "Snippet"},
|
|
{"url": "mailto:invalid@example.org", "title": "Invalid"},
|
|
]
|
|
},
|
|
)
|
|
)
|
|
)
|
|
try:
|
|
results = await provider.search("test")
|
|
assert len(results) == 1
|
|
assert results[0].provider == "searxng"
|
|
assert results[0].url == "https://example.org/a"
|
|
assert results[0].snippet == "Snippet"
|
|
assert provider._client is not None
|
|
finally:
|
|
await provider._client.aclose()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_orchestrator_registers_searxng_provider() -> None:
|
|
orchestrator = ResearchOrchestrator(_config(), uuid4(), "test")
|
|
provider = orchestrator._get_multi_search().get_provider("searxng")
|
|
assert isinstance(provider, SearXNGProvider)
|
|
|
|
|
|
def test_searching_accepts_normalized_result_with_empty_snippet() -> None:
|
|
"""A Pydantic result must not be treated as a mapping during conversion."""
|
|
async def run() -> None:
|
|
orchestrator = ResearchOrchestrator(_config(), uuid4(), "test")
|
|
result = NormalizedResult.from_raw(
|
|
provider_name="searxng",
|
|
title="Example",
|
|
url="https://example.org/article",
|
|
snippet="",
|
|
)
|
|
multi_search = Mock()
|
|
multi_search.enabled_providers.return_value = ["searxng"]
|
|
multi_search.search = AsyncMock(return_value=[result])
|
|
orchestrator._multi_search = multi_search
|
|
|
|
response = await orchestrator._step_searching()
|
|
|
|
assert response["success"] is True
|
|
assert response["url_count"] == 1
|
|
assert orchestrator._search_results == [{
|
|
"url": "https://example.org/article",
|
|
"title": "Example",
|
|
"snippet": "",
|
|
"provider": "searxng",
|
|
}]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_fetching_uses_actual_bytes_and_respects_remaining_source_budget() -> None:
|
|
async def run() -> None:
|
|
budget = HardBudgetConfig(max_sources=2, max_total_download_bytes=1_000)
|
|
orchestrator = ResearchOrchestrator(_config(), uuid4(), "test", budget_config=budget)
|
|
orchestrator._search_results = [
|
|
{"url": "https://example.org/one"},
|
|
{"url": "https://example.org/two"},
|
|
{"url": "https://example.org/three"},
|
|
]
|
|
doc_one = NormalizedDocument.from_text(
|
|
"https://example.org/one", "one", metadata={"download_bytes": 120}
|
|
)
|
|
doc_two = NormalizedDocument.from_text(
|
|
"https://example.org/two", "two", metadata={"download_bytes": 180}
|
|
)
|
|
crawler = Mock()
|
|
crawler.fetch_and_extract_many = AsyncMock(return_value=[doc_one, doc_two])
|
|
orchestrator._crawler = crawler
|
|
|
|
response = await orchestrator._step_fetching()
|
|
|
|
assert response["success"] is True
|
|
crawler.fetch_and_extract_many.assert_awaited_once_with(
|
|
["https://example.org/one", "https://example.org/two"], max_size=500
|
|
)
|
|
usage = orchestrator.budget_tracker.get_usage()
|
|
assert usage["max_sources"]["usage"] == 2
|
|
assert usage["max_total_download_bytes"]["usage"] == 300
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_hard_budget_allows_exactly_the_configured_limit() -> None:
|
|
tracker = BudgetTracker(HardBudgetConfig(max_sources=2))
|
|
tracker.increment_sources(2)
|
|
tracker.check_budget()
|
|
|
|
tracker.increment_sources()
|
|
try:
|
|
tracker.check_budget()
|
|
except BudgetExhaustedError as exc:
|
|
assert exc.exhausted_limits == ["max_sources"]
|
|
else:
|
|
raise AssertionError("Source usage above the limit must exhaust the budget")
|
|
|
|
|
|
def test_time_tracking_records_only_each_new_elapsed_interval() -> None:
|
|
with patch("nsct.orchestration.budget.time.monotonic", side_effect=[100.0, 110.0, 125.0]):
|
|
tracker = BudgetTracker(HardBudgetConfig())
|
|
tracker.record_time_elapsed()
|
|
tracker.record_time_elapsed()
|
|
|
|
usage = tracker.get_usage()
|
|
assert usage["max_research_duration_seconds"]["usage"] == 25.0
|
|
|
|
|
|
def test_depth_time_budgets_allow_local_35b_inference() -> None:
|
|
assert DEPTH_CONFIGS["quick"].max_research_duration_seconds == 1_800
|
|
assert DEPTH_CONFIGS["normal"].max_research_duration_seconds == 3_600
|
|
assert DEPTH_CONFIGS["deep"].max_research_duration_seconds == 7_200
|
|
|
|
|
|
def test_start_does_not_charge_a_phantom_planner_request() -> None:
|
|
async def run() -> None:
|
|
orchestrator = ResearchOrchestrator(_config(), uuid4(), "test")
|
|
|
|
await orchestrator.start()
|
|
|
|
usage = orchestrator.budget_tracker.get_usage()
|
|
assert usage["max_llm_requests"]["usage"] == 0
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_claim_extraction_reserves_synthesis_budget_and_skips_failed_sources() -> None:
|
|
async def run() -> None:
|
|
budget = HardBudgetConfig(max_sources=15, max_llm_requests=15)
|
|
orchestrator = ResearchOrchestrator(_config(), uuid4(), "test", budget_config=budget)
|
|
await orchestrator.start()
|
|
orchestrator._state_machine = StateMachine(ResearchRunState.FETCHING)
|
|
orchestrator._budget_tracker.increment_llm_requests() # completed planning call
|
|
orchestrator._llm_provider = Mock()
|
|
orchestrator._sources = [
|
|
{"id": str(index), "url": f"https://example.org/{index}", "error": ""}
|
|
for index in range(14)
|
|
] + [{"id": "failed", "url": "https://example.org/failed", "error": "timeout"}]
|
|
captured: dict[str, object] = {}
|
|
|
|
class FakeExtractor:
|
|
def __init__(self, **kwargs: object) -> None:
|
|
captured["sources"] = kwargs["sources"]
|
|
|
|
async def extract(self) -> list[object]:
|
|
return []
|
|
|
|
with patch("nsct.stages.stage5_extract_claims.Stage5Extractor", FakeExtractor):
|
|
response = await orchestrator._step_extracting()
|
|
|
|
assert response["success"] is True
|
|
assert len(captured["sources"]) == 13
|
|
usage = orchestrator.budget_tracker.get_usage()
|
|
assert usage["max_llm_requests"]["usage"] == 14
|
|
|
|
# The reserved final request reaches, but does not exceed, the quick limit.
|
|
orchestrator._budget_tracker.increment_llm_requests()
|
|
orchestrator.budget_tracker.check_budget()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_pipeline_executes_extracting_between_fetching_and_analyzing() -> None:
|
|
async def run() -> None:
|
|
orchestrator = ResearchOrchestrator(_config(), uuid4(), "test")
|
|
orchestrator._comparison_data = {}
|
|
for step in ("planning", "searching", "fetching", "extracting", "analyzing", "comparing", "synthesizing"):
|
|
async def handler(target: str = step) -> dict[str, object]:
|
|
orchestrator._transition_to(target)
|
|
return {"success": True, "data": {}}
|
|
|
|
setattr(orchestrator, f"_step_{step}", AsyncMock(side_effect=handler))
|
|
|
|
result = await orchestrator.run()
|
|
|
|
assert result["success"] is True
|
|
assert orchestrator._step_extracting.await_count == 1
|
|
assert orchestrator.state is ResearchRunState.COMPLETED
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_no_search_provider_marks_run_failed_with_reason() -> None:
|
|
async def run() -> None:
|
|
orchestrator = ResearchOrchestrator(_config(searxng_base_url=None), uuid4(), "test")
|
|
await orchestrator.start()
|
|
orchestrator._step_planning = AsyncMock(return_value={"success": True, "data": {}})
|
|
|
|
result = await orchestrator.run()
|
|
|
|
assert result["success"] is False
|
|
assert result["failed_step"] == "searching"
|
|
assert "No search provider" in result["error"]
|
|
assert orchestrator.state is ResearchRunState.FAILED
|
|
|
|
asyncio.run(run())
|