mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-09-18 06:02:21 +02:00
140 lines
4.1 KiB
Python
140 lines
4.1 KiB
Python
import asyncio
|
|
import json
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from src.deep_research import DeepResearcher
|
|
from src.research_navigator import ResearchPage
|
|
|
|
|
|
class _ControlledResearcher(DeepResearcher):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(
|
|
llm_endpoint="http://local.test/v1/chat/completions",
|
|
llm_model="local-model",
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
self.active = 0
|
|
self.max_active = 0
|
|
|
|
async def _search(self, query):
|
|
return [
|
|
{"url": f"https://example.test/{query}/{i}", "title": f"{query}-{i}"}
|
|
for i in range(4)
|
|
]
|
|
|
|
async def _fetch_and_extract(self, url, question, title):
|
|
self.active += 1
|
|
self.max_active = max(self.max_active, self.active)
|
|
await asyncio.sleep(0.01)
|
|
self.active -= 1
|
|
return {"url": url, "title": title, "summary": "ok"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_and_extract_respects_extraction_concurrency():
|
|
researcher = _ControlledResearcher(extraction_concurrency=2, max_urls_per_round=4)
|
|
researcher._start_time = time.time()
|
|
|
|
findings = await researcher._search_and_extract(["a", "b"], "question")
|
|
|
|
assert len(findings) == 8
|
|
assert researcher.max_active == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_and_extract_tracks_all_urls_selected_for_analysis():
|
|
researcher = _ControlledResearcher(extraction_concurrency=2, max_urls_per_round=2)
|
|
researcher._start_time = time.time()
|
|
|
|
findings = await researcher._search_and_extract(["a"], "question")
|
|
|
|
assert len(findings) == 2
|
|
assert [
|
|
{"url": item["url"], "title": item["title"]}
|
|
for item in researcher.analyzed_urls
|
|
] == [
|
|
{"url": "https://example.test/a/0", "title": "a-0"},
|
|
{"url": "https://example.test/a/1", "title": "a-1"},
|
|
]
|
|
assert all(item["requested_by"] == "web_search" for item in researcher.analyzed_urls)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_and_extract_uses_configured_timeout(monkeypatch):
|
|
captured = {}
|
|
|
|
researcher = DeepResearcher(
|
|
llm_endpoint="http://local.test/v1/chat/completions",
|
|
llm_model="local-model",
|
|
extraction_timeout=123,
|
|
)
|
|
|
|
async def fake_fetch(url, timeout=10):
|
|
return ResearchPage(
|
|
url=url,
|
|
title="Page",
|
|
content="useful page content",
|
|
success=True,
|
|
retrieval="fetch",
|
|
)
|
|
|
|
researcher.navigator.fetch = fake_fetch
|
|
|
|
async def fake_llm(messages, temperature=0.3, max_tokens=4096, timeout=60):
|
|
captured["timeout"] = timeout
|
|
return json.dumps({
|
|
"rational": "relevant",
|
|
"evidence": "evidence",
|
|
"summary": "useful page content",
|
|
})
|
|
|
|
researcher._llm = fake_llm
|
|
|
|
result = await researcher._fetch_and_extract("https://example.test", "question", "Title")
|
|
|
|
assert result["summary"] == "useful page content"
|
|
assert captured["timeout"] == 123
|
|
|
|
|
|
def test_extraction_timeout_allows_long_local_model_runs():
|
|
researcher = DeepResearcher(
|
|
llm_endpoint="http://local.test/v1/chat/completions",
|
|
llm_model="local-model",
|
|
extraction_timeout=1800,
|
|
)
|
|
|
|
assert researcher.extraction_timeout == 1800
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_planning_and_query_generation_use_configured_timeouts():
|
|
researcher = DeepResearcher(
|
|
llm_endpoint="http://local.test/v1/chat/completions",
|
|
llm_model="local-model",
|
|
planning_timeout=234,
|
|
query_timeout=345,
|
|
)
|
|
captured = []
|
|
|
|
async def fake_llm(messages, temperature=0.3, max_tokens=4096, timeout=60):
|
|
captured.append(timeout)
|
|
if max_tokens == 1024:
|
|
return json.dumps({
|
|
"sub_questions": ["one"],
|
|
"key_topics": ["topic"],
|
|
"success_criteria": "complete",
|
|
})
|
|
return json.dumps(["query one", "query two"])
|
|
|
|
researcher._llm = fake_llm
|
|
|
|
plan = await researcher._create_plan("question")
|
|
queries = await researcher._generate_queries("question", "", 1)
|
|
|
|
assert "Sub-questions: one" in plan
|
|
assert queries == ["query one", "query two"]
|
|
assert captured == [234, 345]
|