From 3c658e46660f0b1c0ff18286427a401a487f11d8 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 05:46:14 +0000 Subject: [PATCH] Studio: fit Deep Research synthesis evidence to loaded context, add opt-in web grounding Size the synthesis evidence budget to the loaded model context so the prompt is not silently truncated on small contexts. When the evidence overflowed the window the report degenerated (it echoed the evidence tail instead of writing); the budget now reserves tokens for the prompt scaffolding and converts the remainder to chars, keeping the full cap when the context is unknown. Add opt-in web grounding for auto-read: read the top search results, ingest them into an ephemeral RAG scope, hybrid-retrieve the passages most relevant to the question with the existing knowledge-base retriever, and fold those chunks into the step evidence. The scope is per call and deleted afterwards, so a user's knowledge base is never touched. Off by default; enable with UNSLOTH_RESEARCH_AUTO_SCRAPE=1. Gated per run by budgets["maxAutoScrape"], so runs created without it keep legacy snippet-only behavior, and grounding is skipped when the loaded context is too small for the prompt. Add tests for the adaptive evidence budget, scraped-text cleaning, the ephemeral web-RAG retrieval and scope cleanup, and the auto-read evidence path. --- studio/backend/core/rag/web_rank.py | 129 +++++++ studio/backend/core/research_runs.py | 238 +++++++++++- studio/backend/routes/research_runs.py | 8 + .../tests/test_research_runs_storage.py | 357 +++++++++++++++++- studio/backend/tests/test_web_rank.py | 116 ++++++ 5 files changed, 839 insertions(+), 9 deletions(-) create mode 100644 studio/backend/core/rag/web_rank.py create mode 100644 studio/backend/tests/test_web_rank.py diff --git a/studio/backend/core/rag/web_rank.py b/studio/backend/core/rag/web_rank.py new file mode 100644 index 0000000000..1808e2b5e9 --- /dev/null +++ b/studio/backend/core/rag/web_rank.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Ephemeral web-RAG for deep research auto-read. + +Deep research auto-reads the top search results so synthesis is grounded in page text rather +than short snippets. Whole pages make a small local model loop on boilerplate, so the scraped +pages go through the *same* retrieval pipeline the knowledge base uses and only the most +relevant passages are folded into the evidence. + +Nothing here re-implements chunking, embedding, retrieval, ranking, or rendering; it wires +Studio's existing KB components (``chunk_pages``, ``embeddings.encode``, ``store.add_chunks``, +``retrieval.retrieve_hybrid``, ``retrieval.filter_min_score``, ``tool._format``) to the live +scrape. The only difference from a persisted KB is the corpus: pages are ingested under a +unique throwaway scope deleted in a ``finally`` block, so an auto-read never pollutes a user's +knowledge base, exactly like Studio's per-thread attachment RAG on the same store. +""" + +from __future__ import annotations + +import hashlib +import uuid + +from loggers import get_logger +from storage import rag_db + +from . import config, embeddings, retrieval, store, tool +from .chunking import chunk_pages +from .parsers import Page + +logger = get_logger(__name__) + + +def _fit_to_budget(hits, rows, char_budget): + """Keep the best (already ranked) hits whose cumulative chunk text fits ``char_budget``, + always keeping at least the top hit so a single long passage is not dropped whole.""" + if char_budget is None: + return hits + kept = [] + used = 0 + for hit in hits: + row = rows.get(hit.chunk_id) + text = (row["text"] if row else "") or "" + if kept and used + len(text) > char_budget: + break + kept.append(hit) + used += len(text) + return kept + + +def retrieve_web_chunks( + pages: list[dict], + query: str, + *, + top_n: int, + min_score: float, + char_budget: int | None = None, + max_tokens: int | None = None, + overlap: int | None = None, + model_name: str | None = None, +) -> tuple[str, list[dict]]: + """Ingest scraped pages into an ephemeral RAG scope, hybrid-retrieve the passages most + relevant to ``query``, and return ``(rendered_chunks, sources)`` using Studio's KB + formatter. + + ``pages`` is a list of dicts with ``text`` (required) and optional ``title`` / ``url`` + (``title`` becomes the ````). Returns ``("", [])`` when there is nothing + usable or RAG is unavailable, so the caller can fall back to snippet evidence. The scope + is always deleted before returning, so nothing is left in the store.""" + query = (query or "").strip() + if not query or top_n <= 0 or not pages or not rag_db.RAG_AVAILABLE: + return "", [] + model = model_name or config.effective_embedding_model() + max_tokens = max_tokens or config.CHUNK_TOKENS + overlap = config.CHUNK_OVERLAP if overlap is None else overlap + count = embeddings.token_counter(model) + + conn = rag_db.get_connection() + scope = f"research_scrape_{uuid.uuid4().hex}" + doc_ids: list[str] = [] + try: + for page in pages: + text = str(page.get("text") or "").strip() + if not text: + continue + source = str(page.get("title") or page.get("url") or "web").strip() or "web" + chunks = chunk_pages( + [Page(text = text, page_number = None, char_count = len(text))], + max_tokens = max_tokens, + overlap = overlap, + count = count, + ) + if not chunks: + continue + vectors = embeddings.encode( + [chunk.text for chunk in chunks], model_name = model, normalize = True + ) + doc_id = store.create_document( + conn, + scope = scope, + filename = source, + sha256 = hashlib.sha256(text.encode("utf-8", "ignore")).hexdigest(), + status = "ready", + embedding_model = model, + ) + doc_ids.append(doc_id) + store.add_chunks(conn, scope, doc_id, chunks, vectors) + + if not doc_ids: + return "", [] + hits = retrieval.retrieve_hybrid( + conn, scope, query, k = top_n, model_name = model, mode = "hybrid" + ) + hits = retrieval.filter_min_score(hits, min_score) + if not hits: + return "", [] + rows = store.chunks_by_id(conn, [hit.chunk_id for hit in hits]) + hits = _fit_to_budget(hits, rows, char_budget) + return tool._format(rows, hits) + except Exception: + logger.warning("research.web_rank_failed", exc_info = True) + return "", [] + finally: + for doc_id in doc_ids: + try: + store.delete_document(conn, doc_id) + except Exception: + logger.warning("research.web_rank_cleanup_failed doc_id=%s", doc_id) + conn.close() diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index 01dd55e367..3b369ab6c5 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio import ipaddress import json +import os import re import sqlite3 import threading @@ -81,6 +82,78 @@ _MAX_ERROR_CHARS = 500 _MAX_CONTEXT_CHARS = 12_000 _MAX_CONTEXT_MESSAGE_CHARS = 4_000 _MAX_SYNTHESIS_EVIDENCE_CHARS = 32_000 +# The synthesis prompt must fit the loaded context or it is silently truncated and the report +# degenerates (echoes the evidence tail). Studio defaults context to 2048 tokens, far below the +# cap above, so the evidence budget adapts to the loaded context: reserve tokens for the prompt +# scaffolding, then convert the remainder to chars. Unknown context keeps the full cap. +_MIN_SYNTHESIS_EVIDENCE_CHARS = 1_500 +_SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN = 3.0 +_SYNTHESIS_CONTEXT_RESERVE_TOKENS = 2_048 +# Below this loaded context the prompt scaffolding alone fills the window and the grounded +# report degenerates, so grounding is skipped (snippet-only) for smaller loads. +_AUTO_SCRAPE_MIN_CONTEXT_TOKENS = 8_192 +# Optionally read the top search results so synthesis is grounded in page text, not just +# snippets: each scraped page is ingested into an ephemeral RAG scope (deleted after, so a +# user's knowledge base is untouched), the passages most relevant to the question are +# hybrid-retrieved reusing the KB retriever, and the resulting blocks replace the raw +# search text (staying under the existing 12k per-note cap). OFF by default, opt in via +# UNSLOTH_RESEARCH_AUTO_SCRAPE=1: benchmarking showed no reliable factoid-accuracy gain over +# snippets on a local model (snippets usually already carry the fact) while adding latency. +# Gated per run by budgets["maxAutoScrape"] (absent/0 means no scrape, so existing runs keep +# legacy behavior). Safe only with the context gate in _research and the adaptive budget in +# _synthesis_evidence_budget; without them, denser evidence overflows a small context. +_AUTO_SCRAPE_TOP_K = 3 +_AUTO_SCRAPE_TOTAL_CHARS = 6_000 +_WEB_RAG_TOP_N = 6 +_WEB_RAG_MIN_SCORE = 0.30 + + +def _auto_scrape_default() -> int: + """Server default for ``budgets["maxAutoScrape"]``: 0 (off) unless + ``UNSLOTH_RESEARCH_AUTO_SCRAPE`` enables it (``1``/``true`` -> ``_AUTO_SCRAPE_TOP_K``, or an + explicit count clamped to ``[0, _AUTO_SCRAPE_TOP_K]``).""" + raw = os.environ.get("UNSLOTH_RESEARCH_AUTO_SCRAPE", "").strip().lower() + if not raw: + return 0 + if raw in ("0", "false", "no", "off"): + return 0 + if raw in ("1", "true", "yes", "on"): + return _AUTO_SCRAPE_TOP_K + try: + return max(0, min(int(raw), _AUTO_SCRAPE_TOP_K)) + except ValueError: + return 0 + + +# Nav menus, language sidebars, and percent-encoded link lists are not evidence and derail +# retrieval; drop link-dominated and encoded-URL lines. +_MD_LINK = re.compile(r"\[([^\]]*)\]\([^)]*\)") +_PERCENT_ESCAPE = re.compile(r"%[0-9A-Fa-f]{2}") +_LIST_PREFIX = re.compile(r"^(?:[\*\-\+•]|\d+[.)])\s") +_BLANK_RUN = re.compile(r"\n{3,}") +# Bare tracking/redirect URLs arrive as one unbroken token (prose never has an 80-char word); +# not evidence, and a small model will latch onto and echo it. +_LONG_TOKEN = re.compile(r"\S{80,}") + + +def _clean_scraped_text(text: str) -> str: + kept: list[str] = [] + for line in text.splitlines(): + stripped = line.strip() + if not stripped: + kept.append("") + continue + if len(_PERCENT_ESCAPE.findall(stripped)) >= 4: + continue + if _LONG_TOKEN.search(stripped): + continue + prose = _MD_LINK.sub(r"\1", stripped).strip() + if "](" in stripped and ( + _LIST_PREFIX.match(stripped) or len(prose) <= max(30, len(stripped) // 3) + ): + continue + kept.append(line) + return _BLANK_RUN.sub("\n\n", "\n".join(kept)).strip() _REPORT_SYSTEM_PROMPT = """You are writing a rigorous, self-contained research report. @@ -370,13 +443,42 @@ def _research_question_context(thread_id: str, user_message_id: str) -> tuple[st return question, json.dumps(turns, ensure_ascii = False) -def _bounded_synthesis_evidence(notes: list[str]) -> str: +def _loaded_context_length() -> int | None: + """Best-effort read of the active model's context window in tokens, or None if unknown.""" + try: + from core.inference.inference import get_inference_backend + + backend = get_inference_backend() + name = getattr(backend, "active_model_name", None) + if name: + ctx = (getattr(backend, "models", {}).get(name) or {}).get("context_length") + if isinstance(ctx, int) and not isinstance(ctx, bool) and ctx > 0: + return ctx + except Exception: + logger.debug("research.context_probe_failed", exc_info = True) + return None + + +def _synthesis_evidence_budget() -> int: + """Char budget for synthesis evidence, sized to fit the loaded context (falls back to the + full cap when the context is unknown).""" + ctx = _loaded_context_length() + if not ctx: + return _MAX_SYNTHESIS_EVIDENCE_CHARS + usable_tokens = max(0, ctx - _SYNTHESIS_CONTEXT_RESERVE_TOKENS) + budget = int(usable_tokens * _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) + return max(_MIN_SYNTHESIS_EVIDENCE_CHARS, min(budget, _MAX_SYNTHESIS_EVIDENCE_CHARS)) + + +def _bounded_synthesis_evidence( + notes: list[str], max_chars: int = _MAX_SYNTHESIS_EVIDENCE_CHARS +) -> str: if not notes: return "(none)" separator = "\n\n" per_note = max( - 1000, - (_MAX_SYNTHESIS_EVIDENCE_CHARS - len(separator) * (len(notes) - 1)) // len(notes), + min(1000, max_chars), + (max_chars - len(separator) * (len(notes) - 1)) // len(notes), ) bounded = [] for note in notes: @@ -384,7 +486,7 @@ def _bounded_synthesis_evidence(notes: list[str]) -> str: bounded.append(note) else: bounded.append(note[: per_note - 24].rstrip() + "\n[Evidence truncated]") - return separator.join(bounded)[:_MAX_SYNTHESIS_EVIDENCE_CHARS] + return separator.join(bounded)[:max_chars] def _parse_json_object(text: str) -> dict: @@ -749,6 +851,87 @@ class ResearchSupervisor: if self._cancel_event(run_id).is_set(): raise RunCancelled() + async def _auto_scrape_sources( + self, + run: dict, + question: str, + step_sources: list[dict], + fetched_urls: set[str], + *, + tool_timeout: int, + website_policy: dict | None, + ) -> tuple[str, list[str]]: + """Concurrently read the top few of this step's accepted source URLs, rank their + content against the research question with the knowledge-base embedding model, and + return the most relevant chunks as ```` evidence plus the URLs actually read. + URLs are already access checked and deduplicated by the caller, so no new sources are + created. Failures, timeouts, unreadable pages, and low-relevance chunks are dropped; + the caller enforces cancellation.""" + targets = [] + for source in step_sources: + url = str(source.get("url") or "") + if url and url not in fetched_urls: + targets.append(source) + if len(targets) >= _AUTO_SCRAPE_TOP_K: + break + if not targets: + return "", [] + cancel_event = self._cancel_event(run["id"]) + results = await asyncio.gather( + *( + asyncio.to_thread( + execute_tool, + "web_search", + {"url": source["url"]}, + cancel_event = cancel_event, + timeout = tool_timeout, + website_policy = website_policy, + ) + for source in targets + ), + return_exceptions = True, + ) + pages = [] + fetched = [] + for source, result in zip(targets, results): + if isinstance(result, BaseException) or not isinstance(result, str): + continue + body = strip_result_for_model(result) + if is_tool_error(body): + continue + body = _clean_scraped_text(body) + if not body: + continue + fetched.append(source["url"]) + pages.append( + { + "text": body, + "title": source.get("title") or source["url"], + "url": source["url"], + } + ) + if not pages: + return "", [] + # Reuse Studio's knowledge-base RAG pipeline (ingest -> hybrid retrieve -> + # render) over an ephemeral scope; runs off the event loop since embedding and the + # sqlite/vec index work are CPU/GPU bound. + from core.rag import web_rank + + section, _sources = await asyncio.to_thread( + web_rank.retrieve_web_chunks, + pages, + question, + top_n = _WEB_RAG_TOP_N, + min_score = _WEB_RAG_MIN_SCORE, + char_budget = _AUTO_SCRAPE_TOTAL_CHARS, + ) + if not section: + return "", [] + return ( + "Relevant passages retrieved from the top results (already read):\n\n" + section, + fetched, + ) + async def _check_worker_write(self, run_id: str, written: bool) -> None: if written: return @@ -1281,6 +1464,20 @@ class ResearchSupervisor: max_steps = int(budgets["maxSteps"]) max_sources = int(budgets["maxSources"]) tool_timeout = int(budgets["toolTimeoutSeconds"]) + # Absent for runs created before auto-scrape: default 0 keeps their behavior unchanged. + max_auto_scrape = int(budgets.get("maxAutoScrape", 0)) + # Grounding needs the synthesis prompt to fit the loaded context; on a tiny context the + # prompt overhead alone fills the window and the report degenerates, so fall back to + # snippet-only when the context is too small. + if max_auto_scrape > 0: + loaded_ctx = _loaded_context_length() + if loaded_ctx is not None and loaded_ctx < _AUTO_SCRAPE_MIN_CONTEXT_TOKENS: + logger.info( + "research.auto_scrape_disabled_small_context run_id=%s context=%s", + run["id"], + loaded_ctx, + ) + max_auto_scrape = 0 website_policy = run["config"].get("websitePolicy") policy_prompt = website_policy_prompt(website_policy) notes: list[str] = [] @@ -1571,6 +1768,29 @@ class ResearchSupervisor: self.worker_id, ) await self._check_worker_write(run["id"], written) + tool_failed = is_tool_error(result) + step_failed = _research_step_failed(result, rag_sources) + scraped_section = "" + if ( + action["action"] == "search" + and step_sources + and not tool_failed + and max_auto_scrape > 0 + ): + scraped_section, scraped_urls = await self._auto_scrape_sources( + run, + question, + step_sources, + fetched_urls, + tool_timeout = tool_timeout, + website_policy = website_policy, + ) + fetched_urls.update(scraped_urls) + await self._check_active(run["id"]) + if scraped_section: + # Replace raw search text with the retrieved chunks; sources are already + # cataloged above, so nothing citable is lost. + result = scraped_section note = ( f"### {action['title']} ({action['action']})\n" f"Input: {argument}\nResult:\n{result[:12000]}\n\n" @@ -1581,8 +1801,6 @@ class ResearchSupervisor: f"### {action['title']} ({action['action']})\n" f"Input: {argument}\nResult:\n{result[:12000]}" ) - tool_failed = is_tool_error(result) - step_failed = _research_step_failed(result, rag_sources) clean_result = strip_result_for_model(result) step_result = { "action": action["action"], @@ -1590,7 +1808,11 @@ class ResearchSupervisor: "sourceCount": len(step_sources) + len(rag_sources), "sourceUrls": [source["url"] for source in step_sources], "evidenceSources": rag_sources, - **({"excerpt": clean_result[:12000]} if action["action"] == "fetch" else {}), + **( + {"excerpt": clean_result[:12000]} + if action["action"] == "fetch" or scraped_section + else {} + ), **({"error": clean_result[:500]} if tool_failed else {}), } await self._check_active(run["id"]) @@ -1633,7 +1855,7 @@ class ResearchSupervisor: f" Chunk ID: {source.get('chunkId') or '(unknown)'}" for index, source in enumerate(document_sources, 1) ) - evidence_text = _bounded_synthesis_evidence(notes) + evidence_text = _bounded_synthesis_evidence(notes, _synthesis_evidence_budget()) report, synthesis_reasoning, synthesis_finish_reason = await self._stream_completion( run, [ diff --git a/studio/backend/routes/research_runs.py b/studio/backend/routes/research_runs.py index 4c5fd6591b..7faaf4a328 100644 --- a/studio/backend/routes/research_runs.py +++ b/studio/backend/routes/research_runs.py @@ -252,6 +252,14 @@ def _sanitize_config(payload: CreateResearchRun, thread: dict) -> dict: raise HTTPException( status_code = 400, detail = f"{key} must be between {minimum} and {maximum}" ) + # Server-controlled, not client tunable. OFF by default; opt in via + # UNSLOTH_RESEARCH_AUTO_SCRAPE=1. Injected only when enabled, so a default run's budgets stay + # byte-identical to legacy. + from core.research_runs import _auto_scrape_default + + _auto_scrape = _auto_scrape_default() + if _auto_scrape > 0: + budgets["maxAutoScrape"] = _auto_scrape try: website_policy = normalize_website_policy(payload.websitePolicy) except ValueError as exc: diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index b7794f5e4e..2cf9bc2e91 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -55,6 +55,7 @@ def _create( user_message_id = "user-1", rag_scope = None, instructions = "", + budgets = None, ): return research_db.create_run( run_id = run_id, @@ -67,7 +68,7 @@ def _create( "inferenceRequest": {"model": "local-model"}, "ragScope": rag_scope, "instructions": instructions, - "budgets": { + "budgets": budgets or { "maxSteps": 5, "maxSources": 15, "modelTimeoutSeconds": 30, @@ -178,6 +179,31 @@ def test_synthesis_evidence_is_bounded_across_all_steps(): assert all(f"### Step {index}" in evidence for index in range(12)) +def test_synthesis_evidence_budget_tracks_loaded_context(monkeypatch): + from core import research_runs as worker + + # Unknown context keeps the full cap (backwards compatible). + monkeypatch.setattr(worker, "_loaded_context_length", lambda: None) + assert worker._synthesis_evidence_budget() == worker._MAX_SYNTHESIS_EVIDENCE_CHARS + + # A small context (Studio's 2048 default) shrinks the budget so evidence fits. + monkeypatch.setattr(worker, "_loaded_context_length", lambda: 2048) + small = worker._synthesis_evidence_budget() + assert worker._MIN_SYNTHESIS_EVIDENCE_CHARS <= small < worker._MAX_SYNTHESIS_EVIDENCE_CHARS + + # A large context uses (and clamps to) the full cap. + monkeypatch.setattr(worker, "_loaded_context_length", lambda: 32768) + assert worker._synthesis_evidence_budget() == worker._MAX_SYNTHESIS_EVIDENCE_CHARS + + +def test_bounded_synthesis_evidence_respects_small_budget(): + from core import research_runs as worker + + notes = ["### Step\n" + "x" * 20_000 for _ in range(6)] + evidence = worker._bounded_synthesis_evidence(notes, 3_072) + assert len(evidence) <= 3_072 + + def test_report_is_recovered_from_substantial_synthesis_reasoning(): from core import research_runs as worker @@ -1027,6 +1053,7 @@ def test_research_budget_defaults_support_long_runs(): {"modelId": "local-model"}, ) + # auto-scrape (page grounding) is off by default, so budgets stay byte-identical to legacy assert config["budgets"] == { "maxSteps": 12, "maxSources": 40, @@ -1275,6 +1302,334 @@ def test_supervisor_planning_and_research_are_durable_with_mocked_io(research_ho ) +_SCRAPE_BUDGETS = { + "maxSteps": 5, + "maxSources": 15, + "modelTimeoutSeconds": 30, + "toolTimeoutSeconds": 10, + "maxAutoScrape": 3, +} + + +def _patch_web_rank(monkeypatch, *, retrieve = None): + """Stub the ephemeral web-RAG so loop-integration tests need no sqlite/vec store: by + default each scraped page renders as one ```` block, mirroring the real + ``retrieve_web_chunks`` output (whose retrieval/ranking is covered in test_web_rank.py).""" + from core.rag import web_rank + + def default_retrieve(pages, query, *, top_n, min_score, char_budget = None, **kwargs): + blocks, sources = [], [] + for i, page in enumerate(pages, 1): + text = page.get("text") or "" + src = page.get("title") or page.get("url") or "web" + blocks.append(f'\n{text}\n') + sources.append({"citationId": i, "text": text}) + rendered = "\n\n".join(blocks) + if char_budget is not None: + rendered = rendered[:char_budget] + return rendered, sources + + monkeypatch.setattr(web_rank, "retrieve_web_chunks", retrieve or default_retrieve) + + +def _bare_supervisor(monkeypatch): + from core import research_runs as worker + + supervisor = worker.ResearchSupervisor( + SimpleNamespace(state = SimpleNamespace(server_port = 1)) + ) + return worker, supervisor + + +def _run_search_then_finish(monkeypatch, fake_tool, *, retrieve = None): + """Drive one search step (which auto-scrapes) followed by finish, and return the + completed run plus the synthesis prompts the model was given.""" + from core import research_runs as worker + + _patch_web_rank(monkeypatch, retrieve = retrieve) + supervisor = worker.ResearchSupervisor( + SimpleNamespace(state = SimpleNamespace(server_port = 1)) + ) + decisions = iter( + ( + json.dumps({"action": "search", "title": "Find", "query": "grounding evidence"}), + json.dumps({"action": "finish", "title": "Enough evidence"}), + ) + ) + synthesis_prompts = [] + report = "# Report\n\nGrounded finding [source](https://a.example.com)." + + async def fake_stream_completion(run, messages, *, json_mode = False, report_progress = True, **kwargs): + system = messages[0]["content"] + if "rigorous web research plan" in system: + return json.dumps(_plan()), "planned", "stop" + if "iterative research process" in system: + return next(decisions), "decided", "stop" + synthesis_prompts.append(messages[1]["content"]) + research_db.set_report_progress(run["id"], report) + return report, "synthesized", "stop" + + monkeypatch.setattr(supervisor, "_stream_completion", fake_stream_completion) + monkeypatch.setattr(worker, "execute_tool", fake_tool) + + asyncio.run(supervisor._process(research_db.claim_next(supervisor.worker_id))) + planned = research_db.get_run("run-1") + research_db.approve("run-1", planned["planRevision"], planned["planHash"]) + asyncio.run(supervisor._process(research_db.claim_next(supervisor.worker_id))) + return research_db.get_run("run-1"), synthesis_prompts + + +def _two_source_search(): + return ( + "Title: Alpha\nURL: https://a.example.com\nSnippet: alpha snippet.\n\n---\n\n" + "Title: Beta\nURL: https://b.example.com\nSnippet: beta snippet." + ) + + +def test_auto_scrape_retrieves_page_chunks_into_synthesis_evidence(research_home, monkeypatch): + _create(budgets = _SCRAPE_BUDGETS) + url_calls = [] + + def fake_tool(name, arguments, *args, **kwargs): + url = arguments.get("url") + if url: + url_calls.append(url) + return {"https://a.example.com": "ALPHA_PAGE_BODY", "https://b.example.com": "BETA_PAGE_BODY"}[url] + return _two_source_search() + + completed, synthesis_prompts = _run_search_then_finish(monkeypatch, fake_tool) + + assert completed["status"] == "completed" + assert sorted(url_calls) == ["https://a.example.com", "https://b.example.com"] + assert synthesis_prompts, "synthesis must have run" + # the retrieved page chunks reach synthesis, rendered in the format + assert " retrieve -> render reuse chain is exercised end to end) with a fake +deterministic embedding so no model is downloaded. They also assert the ephemeral scope is +deleted, i.e. an auto-read leaves nothing behind in the store.""" + +import numpy as np +import pytest + +from core.rag import web_rank + + +@pytest.fixture +def rag_home(tmp_path, monkeypatch): + """Point rag.db at a throwaway file and rebuild its schema there.""" + from storage import rag_db + + db_file = tmp_path / "rag.db" + monkeypatch.setattr(rag_db, "rag_db_path", lambda: db_file) + monkeypatch.setattr(rag_db, "_schema_ready", False, raising = False) + return db_file + + +@pytest.fixture(autouse = True) +def fake_embeddings(monkeypatch): + """Token counter = word count; embedding = 3-d bag over 'lora'/'license' (+ tiny bias), + so relevance is deterministic and independent of any downloaded model.""" + from core.rag import embeddings as rag_embeddings + + monkeypatch.setattr( + rag_embeddings, + "token_counter", + lambda model_name = None: (lambda text: max(1, len(text.split()))), + ) + + def encode(texts, *, model_name = None, normalize = True): + rows = [] + for text in texts: + low = text.lower() + vec = np.array( + [float(low.count("lora")), float(low.count("license")), 0.001], + dtype = "float32", + ) + norm = np.linalg.norm(vec) + rows.append(vec / norm if (normalize and norm) else vec) + return np.stack(rows) + + monkeypatch.setattr(rag_embeddings, "encode", encode) + + +def _scope_rows(db_file): + """Count leftover ephemeral documents/chunks in the store.""" + import sqlite3 + + conn = sqlite3.connect(str(db_file)) + try: + docs = conn.execute( + "SELECT count(*) FROM documents WHERE scope LIKE 'research_scrape_%'" + ).fetchone()[0] + chunks = conn.execute( + "SELECT count(*) FROM chunks WHERE scope LIKE 'research_scrape_%'" + ).fetchone()[0] + return docs, chunks + finally: + conn.close() + + +def test_retrieves_relevant_passages_as_chunks(rag_home): + pages = [ + {"text": "LoRA is a low-rank adapter method for fine tuning.", "title": "LoRA", "url": "https://a"}, + {"text": "The Apache license governs redistribution terms.", "title": "License", "url": "https://b"}, + ] + rendered, sources = web_rank.retrieve_web_chunks(pages, "what is lora", top_n = 5, min_score = 0.0) + + assert " several ~500-word chunks; a tight budget keeps a bounded subset. + pages = [{"text": " ".join(["lora"] * 2000), "url": "https://a"}] + full, _ = web_rank.retrieve_web_chunks(pages, "lora", top_n = 10, min_score = 0.0) + capped, _ = web_rank.retrieve_web_chunks( + pages, "lora", top_n = 10, min_score = 0.0, char_budget = 3000 + ) + assert full.count("= 2 + assert 1 <= capped.count("