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("