diff --git a/src/fabledassistant/services/tools/web.py b/src/fabledassistant/services/tools/web.py index f95c268..75cfcac 100644 --- a/src/fabledassistant/services/tools/web.py +++ b/src/fabledassistant/services/tools/web.py @@ -23,57 +23,91 @@ logger = logging.getLogger(__name__) required=["query"], ) async def lookup_tool(*, user_id, arguments, **_ctx): + import asyncio + from fabledassistant.config import Config from fabledassistant.services.wikipedia import wiki_summary - query = arguments.get("query", "") + query = arguments.get("query", "").strip() + if not query: + return {"success": False, "error": "query is required"} - # 1. Try Wikipedia first - wiki = await wiki_summary(query) + searxng_enabled = Config.searxng_enabled() + + async def _searxng_results() -> list[dict]: + if not searxng_enabled: + return [] + from fabledassistant.services.research import _search_searxng + return await _search_searxng(query) or [] + + wiki, search_results = await asyncio.gather( + wiki_summary(query), + _searxng_results(), + ) + + wiki_payload: dict | None = None if wiki: + wiki_payload = { + "title": wiki["title"], + "extract": wiki["extract"], + "url": wiki["url"], + } + thumb_url = wiki.get("thumbnail_url") or "" + if thumb_url: + from fabledassistant.services.images import fetch_and_store_image + record = await fetch_and_store_image( + url=thumb_url, + title=wiki["title"], + source_domain="en.wikipedia.org", + ) + if record: + wiki_payload["image"] = { + "embed": f"![{wiki['title']}](/api/images/{record.id})", + "citation": f"*Source: [Wikipedia]({wiki['url']})*", + } + + web_payload: list[dict] = [] + if search_results: + from fabledassistant.services.rss import _fetch_full_article + + top = [r for r in search_results[:2] if r.get("url")] + contents = await asyncio.gather( + *(_fetch_full_article(r["url"]) for r in top), + return_exceptions=True, + ) + for r, content in zip(top, contents, strict=False): + content_text = content if isinstance(content, str) else "" + web_payload.append({ + "url": r["url"], + "title": r.get("title", r["url"]), + "snippet": r.get("snippet", ""), + "content": content_text[:4000], + }) + + if not wiki_payload and not web_payload: return { "success": True, "type": "lookup", - "source": "wikipedia", - "data": wiki, + "data": { + "query": query, + "message": "No results found. You can answer from your own knowledge.", + }, } - # 2. Fall back to SearXNG + article fetch - if Config.searxng_enabled(): - from fabledassistant.services.research import _search_searxng - from fabledassistant.services.rss import _fetch_full_article - - results = await _search_searxng(query) - if results: - articles: list[dict] = [] - for r in results[:2]: - url = r.get("url", "") - if not url: - continue - content = await _fetch_full_article(url) - articles.append({ - "url": url, - "title": r.get("title", url), - "snippet": r.get("snippet", ""), - "content": (content or "")[:4000], - }) - if articles: - return { - "success": True, - "type": "lookup", - "source": "web", - "data": {"query": query, "results": articles}, - } - - # 3. No sources available + data: dict = { + "query": query, + "wikipedia": wiki_payload, + "web": web_payload, + } + if wiki_payload and wiki_payload.get("image"): + data["image_instructions"] = ( + "If an image is relevant to your reply, embed it by writing the wikipedia.image.embed " + "field verbatim, then the citation field on the next line. Otherwise omit it." + ) return { "success": True, "type": "lookup", - "source": "none", - "data": { - "query": query, - "message": "No results found. You can answer from your own knowledge.", - }, + "data": data, } diff --git a/src/fabledassistant/services/wikipedia.py b/src/fabledassistant/services/wikipedia.py index e705b7c..17185cf 100644 --- a/src/fabledassistant/services/wikipedia.py +++ b/src/fabledassistant/services/wikipedia.py @@ -12,11 +12,21 @@ _TIMEOUT = 5.0 _USER_AGENT = "FabledAssistant/1.0 (https://fabledsword.com)" +def _extract_thumbnail(data: dict) -> str: + """Prefer the 320px thumbnail (kinder to cache/bandwidth) over originalimage.""" + thumb = data.get("thumbnail", {}).get("source", "") + if thumb: + return thumb + return data.get("originalimage", {}).get("source", "") + + async def wiki_summary(query: str) -> dict | None: """Fetch a Wikipedia summary for a given title/query. Does a direct title lookup via the REST v1 summary endpoint. - Returns a dict with title, extract, and url on success; None on any failure. + Returns a dict with title, extract, url, and thumbnail_url on success; + thumbnail_url is an empty string when the article has no image. + Returns None on any failure. """ title = url_quote(query.replace(" ", "_")) url = f"{_SUMMARY_URL}/{title}" @@ -37,6 +47,7 @@ async def wiki_summary(query: str) -> dict | None: "title": data.get("title", query), "extract": extract, "url": data.get("content_urls", {}).get("desktop", {}).get("page", ""), + "thumbnail_url": _extract_thumbnail(data), } except Exception as exc: logger.debug("wiki_summary failed for %r: %s", query, exc) @@ -91,6 +102,7 @@ async def wiki_search(query: str, limit: int = 3) -> list[dict]: "url": sdata.get("content_urls", {}) .get("desktop", {}) .get("page", ""), + "thumbnail_url": _extract_thumbnail(sdata), } ) except Exception as exc: diff --git a/tests/test_lookup_tool.py b/tests/test_lookup_tool.py index 34ad27e..f53e73d 100644 --- a/tests/test_lookup_tool.py +++ b/tests/test_lookup_tool.py @@ -1,5 +1,7 @@ import pytest -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch + +from fabledassistant.services.tools.web import lookup_tool @pytest.mark.asyncio @@ -8,16 +10,19 @@ async def test_lookup_wikipedia_hit(): "title": "QUIC", "extract": "QUIC is a transport layer protocol.", "url": "https://en.wikipedia.org/wiki/QUIC", + "thumbnail_url": "", } - with patch("fabledassistant.services.wikipedia.wiki_summary", new_callable=AsyncMock, return_value=wiki_data): - from fabledassistant.services.tools.web import lookup_tool + with patch("fabledassistant.services.wikipedia.wiki_summary", new_callable=AsyncMock, return_value=wiki_data), \ + patch("fabledassistant.config.Config") as mock_config: + mock_config.searxng_enabled.return_value = False result = await lookup_tool(user_id=1, arguments={"query": "QUIC"}) assert result["success"] is True assert result["type"] == "lookup" - assert result["source"] == "wikipedia" - assert result["data"]["title"] == "QUIC" + assert result["data"]["wikipedia"]["title"] == "QUIC" + assert result["data"]["web"] == [] + assert "image" not in result["data"]["wikipedia"] @pytest.mark.asyncio @@ -31,12 +36,58 @@ async def test_lookup_wikipedia_miss_searxng_fallback(): patch("fabledassistant.services.research._search_searxng", new_callable=AsyncMock, return_value=searxng_results), \ patch("fabledassistant.services.rss._fetch_full_article", new_callable=AsyncMock, return_value="Full article about QUIC..."): mock_config.searxng_enabled.return_value = True - from fabledassistant.services.tools.web import lookup_tool result = await lookup_tool(user_id=1, arguments={"query": "QUIC"}) assert result["success"] is True - assert result["source"] == "web" - assert result["data"]["results"] + assert result["data"]["wikipedia"] is None + assert len(result["data"]["web"]) == 1 + assert result["data"]["web"][0]["url"] == "https://example.com/quic" + + +@pytest.mark.asyncio +async def test_lookup_parallel_both_sources(): + wiki_data = { + "title": "President of the United States", + "extract": "The president is the head of state.", + "url": "https://en.wikipedia.org/wiki/President_of_the_United_States", + "thumbnail_url": "", + } + searxng_results = [ + {"url": "https://whitehouse.gov", "title": "The White House", "snippet": "Current admin..."}, + ] + + with patch("fabledassistant.services.wikipedia.wiki_summary", new_callable=AsyncMock, return_value=wiki_data), \ + patch("fabledassistant.config.Config") as mock_config, \ + patch("fabledassistant.services.research._search_searxng", new_callable=AsyncMock, return_value=searxng_results), \ + patch("fabledassistant.services.rss._fetch_full_article", new_callable=AsyncMock, return_value="Current president is..."): + mock_config.searxng_enabled.return_value = True + result = await lookup_tool(user_id=1, arguments={"query": "president of the united states"}) + + assert result["success"] is True + assert result["data"]["wikipedia"]["title"] == "President of the United States" + assert len(result["data"]["web"]) == 1 + + +@pytest.mark.asyncio +async def test_lookup_wikipedia_thumbnail_cached(): + wiki_data = { + "title": "Hedgehog", + "extract": "Hedgehogs are spiny mammals.", + "url": "https://en.wikipedia.org/wiki/Hedgehog", + "thumbnail_url": "https://upload.wikimedia.org/.../Hedgehog.jpg", + } + mock_record = MagicMock() + mock_record.id = 42 + + with patch("fabledassistant.services.wikipedia.wiki_summary", new_callable=AsyncMock, return_value=wiki_data), \ + patch("fabledassistant.config.Config") as mock_config, \ + patch("fabledassistant.services.images.fetch_and_store_image", new_callable=AsyncMock, return_value=mock_record): + mock_config.searxng_enabled.return_value = False + result = await lookup_tool(user_id=1, arguments={"query": "hedgehog"}) + + assert result["data"]["wikipedia"]["image"]["embed"] == "![Hedgehog](/api/images/42)" + assert "Wikipedia" in result["data"]["wikipedia"]["image"]["citation"] + assert "image_instructions" in result["data"] @pytest.mark.asyncio @@ -44,11 +95,10 @@ async def test_lookup_wikipedia_miss_no_searxng(): with patch("fabledassistant.services.wikipedia.wiki_summary", new_callable=AsyncMock, return_value=None), \ patch("fabledassistant.config.Config") as mock_config: mock_config.searxng_enabled.return_value = False - from fabledassistant.services.tools.web import lookup_tool result = await lookup_tool(user_id=1, arguments={"query": "xyznonexistent"}) assert result["success"] is True - assert result["source"] == "none" + assert "message" in result["data"] @pytest.mark.asyncio diff --git a/tests/test_wikipedia.py b/tests/test_wikipedia.py index e82e0b2..6b8f2ac 100644 --- a/tests/test_wikipedia.py +++ b/tests/test_wikipedia.py @@ -37,6 +37,33 @@ async def test_wiki_summary_returns_extract(): assert result["title"] == "Python (programming language)" assert "Python" in result["extract"] assert result["url"].startswith("https://") + assert result["thumbnail_url"] == "" + + +@pytest.mark.asyncio +async def test_wiki_summary_includes_thumbnail(): + """Thumbnail URL is extracted when the article has an image.""" + payload = { + "type": "standard", + "title": "Hedgehog", + "extract": "Hedgehogs are spiny mammals.", + "content_urls": {"desktop": {"page": "https://en.wikipedia.org/wiki/Hedgehog"}}, + "thumbnail": {"source": "https://upload.wikimedia.org/foo-320px.jpg"}, + "originalimage": {"source": "https://upload.wikimedia.org/foo.jpg"}, + } + mock_resp = _make_mock_response(payload) + + with patch("fabledassistant.services.wikipedia.httpx.AsyncClient") as mock_client_cls: + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_client.get = AsyncMock(return_value=mock_resp) + mock_client_cls.return_value = mock_client + + result = await wiki_summary("Hedgehog") + + assert result is not None + assert result["thumbnail_url"] == "https://upload.wikimedia.org/foo-320px.jpg" @pytest.mark.asyncio