From 8db6b4d23095e9b89636cf16349e084948d57b0e Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Fri, 17 Apr 2026 21:59:35 -0400 Subject: [PATCH] feat: add Wikipedia as research pipeline source alongside SearXNG Runs wiki_search in parallel with SearXNG queries; Wikipedia results (which already carry content via their extract field) are merged into the source pool before outline generation, skipping a separate fetch step. Also fixes a pre-existing F811 ruff violation in the test file. Co-Authored-By: Claude Sonnet 4.6 --- src/fabledassistant/services/research.py | 34 +++++++++++++++++--- tests/test_research_pipeline.py | 41 ++++++++++++++++++++++-- 2 files changed, 69 insertions(+), 6 deletions(-) diff --git a/src/fabledassistant/services/research.py b/src/fabledassistant/services/research.py index 8da1eb9..c882e64 100644 --- a/src/fabledassistant/services/research.py +++ b/src/fabledassistant/services/research.py @@ -10,6 +10,7 @@ import httpx from fabledassistant.config import Config from fabledassistant.services.llm import fetch_url_content, generate_completion, stream_chat from fabledassistant.services.notes import create_note, update_note +from fabledassistant.services.wikipedia import wiki_search from fabledassistant.models.note import Note logger = logging.getLogger(__name__) @@ -205,7 +206,7 @@ async def run_research_pipeline( queries = await _generate_sub_queries(topic, model) logger.info("Research: generated %d sub-queries for topic '%s'", len(queries), topic) - # Step 2: Search all queries in parallel (200 ms stagger to avoid hammering SearXNG) + # Step 2: Search all queries in parallel (SearXNG + Wikipedia) async def _search_with_stagger(i: int, query: str) -> tuple[str, list[dict]]: if i > 0: await asyncio.sleep(0.2 * i) @@ -214,13 +215,22 @@ async def run_research_pipeline( logger.info("Research: query '%s' → %d results", query, len(results)) return query, results - search_results = await asyncio.gather( + async def _wiki_for_query(query: str) -> list[dict]: + return await wiki_search(query, limit=1) + + searxng_task = asyncio.gather( *[_search_with_stagger(i, q) for i, q in enumerate(queries)] ) + wiki_task = asyncio.gather( + *[_wiki_for_query(q) for q in queries] + ) + search_results, wiki_results = await asyncio.gather(searxng_task, wiki_task) # Deduplicate URLs across all queries seen_urls: set[str] = set() url_tasks: list[tuple[str, dict, str]] = [] # (url, result_dict, query) + wiki_sources: list[dict] = [] # Wikipedia articles (already have content) + for query, results in search_results: for result in results[:PAGES_PER_QUERY]: url = result.get("url", "") @@ -228,7 +238,21 @@ async def run_research_pipeline( seen_urls.add(url) url_tasks.append((url, result, query)) - # Fetch all unique URLs in parallel + # Add Wikipedia results (they already have content via extract) + for query, wiki_hits in zip(queries, wiki_results): + for hit in wiki_hits: + url = hit.get("url", "") + if url and url not in seen_urls: + seen_urls.add(url) + wiki_sources.append({ + "url": url, + "title": hit["title"], + "query": query, + "snippet": hit["extract"][:200], + "content": hit["extract"], + }) + + # Fetch all unique SearXNG URLs in parallel async def _fetch_source(url: str, result: dict, query: str) -> dict: title = result.get("title", url) _status(f"Reading: {title[:60]}...") @@ -241,10 +265,12 @@ async def run_research_pipeline( "content": content, } - all_sources: list[dict] = list(await asyncio.gather( + fetched_sources: list[dict] = list(await asyncio.gather( *[_fetch_source(url, result, query) for url, result, query in url_tasks] )) + all_sources = wiki_sources + fetched_sources + if not all_sources: raise ValueError(f"No results found for '{topic}'") diff --git a/tests/test_research_pipeline.py b/tests/test_research_pipeline.py index 9c61422..8974d5a 100644 --- a/tests/test_research_pipeline.py +++ b/tests/test_research_pipeline.py @@ -123,8 +123,6 @@ async def test_pipeline_creates_section_notes_and_index(): n.title = title return n - mock_update = AsyncMock() - with patch("fabledassistant.services.research._generate_sub_queries", new_callable=AsyncMock, return_value=["q1"]), \ patch("fabledassistant.services.research._search_searxng", new_callable=AsyncMock, return_value=[{"url": "http://x.com", "title": "X", "snippet": "s"}]), \ patch("fabledassistant.services.research.fetch_url_content", new_callable=AsyncMock, return_value="content"), \ @@ -192,3 +190,42 @@ async def test_pipeline_falls_back_when_all_sections_fail(): result = await run_research_pipeline("test topic", user_id=1, model="test-model") assert result.id == 77 + + +@pytest.mark.asyncio +async def test_pipeline_includes_wikipedia_sources(): + """run_research_pipeline should merge Wikipedia results into the source pool.""" + from unittest.mock import MagicMock + + wiki_results = [{"title": "Wiki Article", "extract": "Wikipedia content about the topic.", "url": "https://en.wikipedia.org/wiki/Topic"}] + + outline = [ + {"title": "Section A", "focus": "Focus A"}, + {"title": "Section B", "focus": "Focus B"}, + ] + + note_id_counter = iter(range(30, 40)) + + def _make_note(user_id, title, body, tags, project_id=None, parent_id=None): + n = MagicMock() + n.id = next(note_id_counter) + n.title = title + return n + + with patch("fabledassistant.services.research._generate_sub_queries", new_callable=AsyncMock, return_value=["q1"]), \ + patch("fabledassistant.services.research._search_searxng", new_callable=AsyncMock, return_value=[{"url": "http://x.com", "title": "X", "snippet": "s"}]), \ + patch("fabledassistant.services.research.wiki_search", new_callable=AsyncMock, return_value=wiki_results), \ + patch("fabledassistant.services.research.fetch_url_content", new_callable=AsyncMock, return_value="content"), \ + patch("fabledassistant.services.research._generate_outline", new_callable=AsyncMock, return_value=outline) as mock_outline, \ + patch("fabledassistant.services.research._synthesize_section", new_callable=AsyncMock, side_effect=lambda t, f, s, m: (t, f"Body for {t}")), \ + patch("fabledassistant.services.research._generate_executive_summary", new_callable=AsyncMock, return_value="Summary."), \ + patch("fabledassistant.services.research.create_note", new_callable=AsyncMock, side_effect=_make_note), \ + patch("fabledassistant.services.research.update_note", new_callable=AsyncMock): + + from fabledassistant.services.research import run_research_pipeline + await run_research_pipeline("test topic", user_id=1, model="test-model") + + # The sources passed to _generate_outline should include the Wikipedia article + sources_arg = mock_outline.call_args[0][1] # second positional arg + source_urls = [s["url"] for s in sources_arg] + assert "https://en.wikipedia.org/wiki/Topic" in source_urls