diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_discover/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_discover/__init__.py new file mode 100644 index 000000000..9ad4346a6 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_discover/__init__.py @@ -0,0 +1 @@ +"""web_discover tool: UI emission for the web.discover capability (05).""" diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_discover/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_discover/emission.py new file mode 100644 index 000000000..4f0592889 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_discover/emission.py @@ -0,0 +1,29 @@ +"""web_discover: ranked-hits card + a result-count terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + if not isinstance(out, dict): + message = str(out) + yield ctx.emit_tool_output_card({"status": "error", "message": message}) + yield ctx.streaming_service.format_terminal_info(message, "error") + return + + hits = out.get("hits") or [] + yield ctx.emit_tool_output_card( + {"status": "completed", "hits": hits, "count": len(hits)} + ) + level = "success" if hits else "info" + yield ctx.streaming_service.format_terminal_info( + f"Found {len(hits)} result(s)", level + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_scrape/__init__.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_scrape/__init__.py new file mode 100644 index 000000000..2282f2277 --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_scrape/__init__.py @@ -0,0 +1 @@ +"""web_scrape tool: UI emission for the web.scrape capability (05).""" diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_scrape/emission.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_scrape/emission.py new file mode 100644 index 000000000..dcc0a494d --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/web_scrape/emission.py @@ -0,0 +1,54 @@ +"""web_scrape: per-page card (content previewed) + a scraped-count terminal line.""" + +from __future__ import annotations + +from collections.abc import Iterator + +from app.tasks.chat.streaming.handlers.tools.emission_context import ( + ToolCompletionEmissionContext, +) + +_PREVIEW_CHARS = 500 + + +def iter_completion_emission_frames( + ctx: ToolCompletionEmissionContext, +) -> Iterator[str]: + out = ctx.tool_output + if not isinstance(out, dict): + message = str(out) + yield ctx.emit_tool_output_card({"status": "error", "message": message}) + yield ctx.streaming_service.format_terminal_info(message, "error") + return + + rows = out.get("rows") or [] + pages = [_page(row) for row in rows] + succeeded = sum(1 for row in rows if row.get("status") == "success") + + yield ctx.emit_tool_output_card( + { + "status": "completed", + "pages": pages, + "succeeded": succeeded, + "total": len(rows), + } + ) + level = "success" if succeeded else "error" + yield ctx.streaming_service.format_terminal_info( + f"Scraped {succeeded}/{len(rows)} page(s)", level + ) + + +def _page(row: dict) -> dict: + """A card-safe view of one row: metadata kept, content bounded to a preview.""" + page: dict = {"url": row.get("url"), "status": row.get("status")} + if row.get("metadata"): + page["metadata"] = row["metadata"] + content = row.get("content") + if content: + page["content_preview"] = ( + content[:_PREVIEW_CHARS] + "…" if len(content) > _PREVIEW_CHARS else content + ) + if row.get("error"): + page["error"] = row["error"] + return page diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_capability_emissions.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_capability_emissions.py new file mode 100644 index 000000000..46a373b8d --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_capability_emissions.py @@ -0,0 +1,131 @@ +"""Emission handlers for the capability chat tools (web_scrape / web_discover).""" + +from __future__ import annotations + +import importlib + + +class _FakeStreaming: + def __init__(self) -> None: + self.terminals: list[tuple[str, str]] = [] + + def format_terminal_info(self, text: str, message_type: str = "info") -> str: + self.terminals.append((text, message_type)) + return f"TERM::{message_type}" + + +class _FakeCtx: + """Minimal stand-in for ToolCompletionEmissionContext (only what handlers use).""" + + def __init__(self, tool_output: object) -> None: + self.tool_output = tool_output + self.cards: list[dict] = [] + self.streaming_service = _FakeStreaming() + + def emit_tool_output_card(self, payload: dict) -> str: + self.cards.append(payload) + return "CARD" + + +def _scrape_frames(ctx: _FakeCtx) -> list[str]: + mod = importlib.import_module( + "app.tasks.chat.streaming.handlers.tools.web_scrape.emission" + ) + return list(mod.iter_completion_emission_frames(ctx)) + + +def _discover_frames(ctx: _FakeCtx) -> list[str]: + mod = importlib.import_module( + "app.tasks.chat.streaming.handlers.tools.web_discover.emission" + ) + return list(mod.iter_completion_emission_frames(ctx)) + + +def test_scrape_card_previews_content_and_summarizes(): + long_body = "x" * 2000 + ctx = _FakeCtx( + { + "rows": [ + { + "url": "https://a.com", + "status": "success", + "content": long_body, + "metadata": {"title": "A"}, + }, + {"url": "https://b.com", "status": "failed", "error": "boom"}, + ] + } + ) + + _scrape_frames(ctx) + + [card] = ctx.cards + assert card["succeeded"] == 1 + assert card["total"] == 2 + page = card["pages"][0] + # full content must not be dumped into the card; a bounded preview is used + assert "content" not in page + assert len(page["content_preview"]) < len(long_body) + assert page["metadata"] == {"title": "A"} + assert card["pages"][1]["error"] == "boom" + + +def test_scrape_terminal_success_when_any_page_succeeds(): + ctx = _FakeCtx( + {"rows": [{"url": "https://a.com", "status": "success", "content": "hi"}]} + ) + + _scrape_frames(ctx) + + assert ctx.streaming_service.terminals[-1][1] == "success" + + +def test_scrape_terminal_error_when_all_pages_fail(): + ctx = _FakeCtx( + {"rows": [{"url": "https://a.com", "status": "failed", "error": "boom"}]} + ) + + _scrape_frames(ctx) + + assert ctx.streaming_service.terminals[-1][1] == "error" + + +def test_scrape_non_dict_output_is_an_error_card(): + ctx = _FakeCtx("Insufficient credit to continue.") + + _scrape_frames(ctx) + + assert ctx.cards[0]["status"] == "error" + assert ctx.streaming_service.terminals[-1][1] == "error" + + +def test_discover_card_lists_hits_and_counts(): + ctx = _FakeCtx( + { + "hits": [ + {"url": "https://a.com", "title": "A", "snippet": "s", "provider": "p"}, + { + "url": "https://b.com", + "title": "B", + "snippet": None, + "provider": "p", + }, + ] + } + ) + + _discover_frames(ctx) + + [card] = ctx.cards + assert card["count"] == 2 + assert card["hits"][0]["url"] == "https://a.com" + assert ctx.streaming_service.terminals[-1][1] == "success" + + +def test_discover_non_dict_output_is_an_error_card(): + ctx = _FakeCtx("Search is not available.") + + _discover_frames(ctx) + + assert ctx.cards[0]["status"] == "error" + assert ctx.streaming_service.terminals[-1][1] == "error"