feat: add DeepSearch multi-agent researcher with async jobs, vector store, and RAG chat

2026-06-14 01:57:17 +00:00
parent 4ae0b0db5d
commit 33d17db79a
8 changed files with 444 additions and 3 deletions
@@ -0,0 +1,55 @@
+# retoor <retoor@molodetz.nl>
+
+import asyncio
+
+from devplacepy.services.deepsearch import chat as chat_module
+from devplacepy.services.deepsearch.chat import DeepsearchChat
+from devplacepy.services.deepsearch.embeddings import local_embed
+from devplacepy.services.deepsearch.store import Chunk, VectorStore
+
+
+def _seed(collection):
+    store = VectorStore(collection)
+    chunks = [
+        Chunk(uid="c0", text="The transistor was invented at Bell Labs.", url="https://a.example", title="A"),
+        Chunk(uid="c1", text="Silicon wafers are used to make chips.", url="https://b.example", title="B"),
+    ]
+    vectors = local_embed([c.text for c in chunks]).vectors
+    store.add(chunks, vectors)
+
+
+def test_answer_is_grounded_and_cited(monkeypatch):
+    collection = "ds_chat_test_grounded"
+    _seed(collection)
+    try:
+        async def fake_embed(texts, api_key, **kwargs):
+            return local_embed(texts)
+
+        async def fake_complete(messages, api_key, **kwargs):
+            return "The transistor was invented at Bell Labs [1]."
+
+        monkeypatch.setattr(chat_module, "embed_texts", fake_embed)
+        monkeypatch.setattr(chat_module, "complete_chat", fake_complete)
+
+        chat = DeepsearchChat(collection, "k")
+        answer = asyncio.run(chat.answer("where was the transistor invented"))
+        assert "Bell Labs" in answer.text
+        assert answer.citations
+        assert answer.citations[0]["url"].startswith("https://")
+    finally:
+        VectorStore(collection).drop()
+
+
+def test_answer_when_no_chunks(monkeypatch):
+    collection = "ds_chat_test_empty"
+    try:
+        async def fake_embed(texts, api_key, **kwargs):
+            return local_embed(texts)
+
+        monkeypatch.setattr(chat_module, "embed_texts", fake_embed)
+        chat = DeepsearchChat(collection, "k")
+        answer = asyncio.run(chat.answer("anything"))
+        assert answer.citations == []
+        assert "did not capture" in answer.text
+    finally:
+        VectorStore(collection).drop()
@@ -0,0 +1,52 @@
+# retoor <retoor@molodetz.nl>
+
+from devplacepy.services.deepsearch.embeddings import local_embed
+from devplacepy.services.deepsearch.store import Chunk, VectorStore
+
+
+def _chunks():
+    texts = [
+        "The transistor was invented at Bell Labs in nineteen forty seven.",
+        "Silicon is the primary semiconductor material used in chips.",
+        "Quantum tunnelling limits how small a transistor can become.",
+    ]
+    chunks = [
+        Chunk(uid=f"c{i}", text=text, url=f"https://s{i}.example", title=f"Source {i}")
+        for i, text in enumerate(texts)
+    ]
+    return chunks
+
+
+def test_add_and_count():
+    store = VectorStore("ds_store_test_count")
+    try:
+        chunks = _chunks()
+        vectors = local_embed([c.text for c in chunks]).vectors
+        store.add(chunks, vectors)
+        assert store.count() == 3
+    finally:
+        store.drop()
+
+
+def test_hybrid_search_returns_relevant_chunk():
+    store = VectorStore("ds_store_test_hybrid")
+    try:
+        chunks = _chunks()
+        vectors = local_embed([c.text for c in chunks]).vectors
+        store.add(chunks, vectors)
+        query = "where was the transistor invented"
+        query_vector = local_embed([query]).vectors[0]
+        results = store.hybrid_search(query, query_vector, top_k=2)
+        assert results
+        assert any("transistor" in r.text.lower() for r in results)
+    finally:
+        store.drop()
+
+
+def test_keyword_scores_rank_match_higher():
+    store = VectorStore("ds_store_test_keyword")
+    chunks = _chunks()
+    scores = store.keyword_scores("silicon semiconductor", chunks)
+    assert scores
+    best = max(scores, key=scores.get)
+    assert "silicon" in next(c.text for c in chunks if c.uid == best).lower()
@@ -0,0 +1,103 @@
+# retoor <retoor@molodetz.nl>
+
+import asyncio
+import json
+import tempfile
+from pathlib import Path
+
+from devplacepy.services.jobs.deepsearch import crawl as crawl_module
+from devplacepy.services.jobs.deepsearch import worker as worker_module
+from devplacepy.services.jobs.deepsearch.crawl import CrawledPage, CrawlOutcome
+from devplacepy.services.deepsearch.embeddings import local_embed
+from devplacepy.services.deepsearch.store import VectorStore
+
+
+def _patch_pipeline(monkeypatch, pages):
+    async def fake_plan(query, api_key):
+        return [query, f"{query} overview"]
+
+    async def fake_search(queries):
+        return [{"url": page.url, "title": page.title, "description": ""} for page in pages]
+
+    async def fake_crawl(candidates, max_pages, emit, is_cached, should_stop):
+        outcome = CrawlOutcome()
+        for page in pages[:max_pages]:
+            emit({"type": "page_loaded", "url": page.url, "done": 1, "total": len(pages)})
+            outcome.pages.append(page)
+        return outcome
+
+    def fake_embed(texts, api_key):
+        return local_embed(texts)
+
+    async def fake_embed_async(texts, api_key, **kwargs):
+        return local_embed(texts)
+
+    async def fake_orchestrate(question, crawled, api_key, emit):
+        from devplacepy.services.jobs.deepsearch.orchestrate import Orchestration
+
+        return Orchestration(
+            summary="A summary.",
+            findings=[{"title": "F", "detail": "D", "confidence": 0.5, "citations": [1]}],
+            gaps=["gap"],
+            confidence=0.6,
+            source_diversity=0.5,
+            score=42,
+        )
+
+    monkeypatch.setattr(worker_module, "plan_queries", fake_plan)
+    monkeypatch.setattr(worker_module, "search_queries", fake_search)
+    monkeypatch.setattr(worker_module, "crawl", fake_crawl)
+    monkeypatch.setattr(worker_module, "embed_texts", fake_embed_async)
+    monkeypatch.setattr(worker_module, "orchestrate", fake_orchestrate)
+
+
+def test_worker_run_produces_report(monkeypatch):
+    pages = [
+        CrawledPage(
+            url="https://example.com/a",
+            title="Page A",
+            text="The transistor was invented at Bell Labs. " * 20,
+            source="httpx",
+            status=200,
+        ),
+        CrawledPage(
+            url="https://other.example/b",
+            title="Page B",
+            text="Semiconductors are made from silicon. " * 20,
+            source="playwright",
+            status=200,
+        ),
+    ]
+    _patch_pipeline(monkeypatch, pages)
+    with tempfile.TemporaryDirectory() as tmp:
+        output_dir = Path(tmp)
+        payload = {
+            "query": "history of the transistor",
+            "max_pages": 5,
+            "depth": 2,
+            "api_key": "k",
+            "collection": "ds_worker_test_one",
+            "cached_hashes": [],
+        }
+        report = asyncio.run(worker_module._run(payload, output_dir))
+        assert report["query"] == "history of the transistor"
+        assert report["page_count"] == 2
+        assert report["chunk_count"] > 0
+        assert report["summary"] == "A summary."
+        assert report["findings"]
+        assert (output_dir / "report.json").is_file()
+        cache = json.loads((output_dir / "url_cache.json").read_text())
+        assert len(cache) == 2
+        VectorStore("ds_worker_test_one").drop()
+
+
+def test_worker_control_cancel_stops(monkeypatch):
+    pages = [
+        CrawledPage(url="https://x.example", title="X", text="content " * 40, source="httpx", status=200)
+    ]
+    _patch_pipeline(monkeypatch, pages)
+    with tempfile.TemporaryDirectory() as tmp:
+        output_dir = Path(tmp)
+        (output_dir / "control.json").write_text(json.dumps({"state": "cancelled"}))
+        should_stop = worker_module._make_stop(output_dir)
+        assert asyncio.run(should_stop()) is True