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.
This commit is contained in:
danielhanchen
parent
2d468916d3
commit
3c658e4666
5 changed files with 839 additions and 9 deletions
129
studio/backend/core/rag/web_rank.py
Normal file
129
studio/backend/core/rag/web_rank.py
Normal file
|
|
@ -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 ``<chunk source>``). 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()
|
||||
|
|
@ -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 <chunk> 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 ``<chunk>`` 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 -> <chunk>
|
||||
# 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,
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 ``<chunk>`` 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'<chunk id="{i}" source="{src}">\n{text}\n</chunk>')
|
||||
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 <chunk> format
|
||||
assert "<chunk" in synthesis_prompts[0]
|
||||
assert "ALPHA_PAGE_BODY" in synthesis_prompts[0]
|
||||
assert "BETA_PAGE_BODY" in synthesis_prompts[0]
|
||||
|
||||
|
||||
def test_auto_scrape_persists_chunk_excerpt_for_resume(research_home, monkeypatch):
|
||||
_create(budgets = _SCRAPE_BUDGETS)
|
||||
|
||||
def fake_tool(name, arguments, *args, **kwargs):
|
||||
url = arguments.get("url")
|
||||
if url:
|
||||
return {"https://a.example.com": "ALPHA_PAGE_BODY", "https://b.example.com": "BETA_PAGE_BODY"}[url]
|
||||
return _two_source_search()
|
||||
|
||||
completed, _ = _run_search_then_finish(monkeypatch, fake_tool)
|
||||
|
||||
search_step = completed["steps"][0]
|
||||
result = search_step["result"]
|
||||
assert result["action"] == "search"
|
||||
assert result["sourceUrls"] == ["https://a.example.com", "https://b.example.com"]
|
||||
assert result["sourceCount"] == 2
|
||||
# the durable excerpt carries the chunks so a resumed run reconstructs the same evidence
|
||||
assert "<chunk" in result["excerpt"]
|
||||
assert "ALPHA_PAGE_BODY" in result["excerpt"]
|
||||
|
||||
|
||||
def test_auto_scrape_ignores_fetch_failures(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 "Error: boom" if url == "https://a.example.com" else "BETA_PAGE_BODY"
|
||||
return _two_source_search()
|
||||
|
||||
completed, synthesis_prompts = _run_search_then_finish(monkeypatch, fake_tool)
|
||||
|
||||
assert completed["status"] == "completed"
|
||||
assert completed["steps"][0]["status"] == "completed"
|
||||
assert len(url_calls) == 2
|
||||
# the failed fetch is never chunked; only the good page's content appears
|
||||
assert "BETA_PAGE_BODY" in synthesis_prompts[0]
|
||||
assert "Error: boom" not in synthesis_prompts[0]
|
||||
|
||||
|
||||
def test_auto_scrape_skipped_for_legacy_config_without_key(research_home, monkeypatch):
|
||||
# Existing/legacy runs persisted no maxAutoScrape; they must never gain scraping on resume
|
||||
# or new steps, regardless of the current server default.
|
||||
_create() # legacy budgets, no maxAutoScrape
|
||||
url_calls = []
|
||||
|
||||
def fake_tool(name, arguments, *args, **kwargs):
|
||||
if arguments.get("url"):
|
||||
url_calls.append(arguments["url"])
|
||||
return "SHOULD_NOT_BE_FETCHED"
|
||||
return _two_source_search()
|
||||
|
||||
completed, synthesis_prompts = _run_search_then_finish(monkeypatch, fake_tool)
|
||||
|
||||
assert completed["status"] == "completed"
|
||||
assert url_calls == []
|
||||
assert "SHOULD_NOT_BE_FETCHED" not in synthesis_prompts[0]
|
||||
assert "excerpt" not in completed["steps"][0]["result"]
|
||||
|
||||
|
||||
def test_auto_scrape_skipped_on_small_context(research_home, monkeypatch):
|
||||
# A context too small for the grounded synthesis prompt would degenerate the report, so
|
||||
# grounding is skipped (snippet-only) even when maxAutoScrape is set.
|
||||
from core import research_runs as worker
|
||||
|
||||
monkeypatch.setattr(worker, "_loaded_context_length", lambda: 2048)
|
||||
_create(budgets = _SCRAPE_BUDGETS)
|
||||
|
||||
def fake_tool(name, arguments, *args, **kwargs):
|
||||
if arguments.get("url"):
|
||||
return "SHOULD_NOT_BE_FETCHED"
|
||||
return _two_source_search()
|
||||
|
||||
completed, synthesis_prompts = _run_search_then_finish(monkeypatch, fake_tool)
|
||||
|
||||
assert completed["status"] == "completed"
|
||||
assert "<chunk" not in synthesis_prompts[0]
|
||||
assert "SHOULD_NOT_BE_FETCHED" not in synthesis_prompts[0]
|
||||
assert "excerpt" not in completed["steps"][0]["result"]
|
||||
|
||||
|
||||
def test_synthesis_pass_runs_at_synthesis_phase(research_home, monkeypatch):
|
||||
# The report pass runs at phase "synthesis" and with default sampling: no repetition
|
||||
# penalty is injected (an aggressive one degenerates small local models into a word-salad).
|
||||
from core import research_runs as worker
|
||||
|
||||
_create(budgets = _SCRAPE_BUDGETS)
|
||||
_patch_web_rank(monkeypatch)
|
||||
supervisor = worker.ResearchSupervisor(
|
||||
SimpleNamespace(state = SimpleNamespace(server_port = 1))
|
||||
)
|
||||
decisions = iter(
|
||||
(
|
||||
json.dumps({"action": "search", "title": "Find", "query": "q"}),
|
||||
json.dumps({"action": "finish", "title": "done"}),
|
||||
)
|
||||
)
|
||||
captured = {}
|
||||
|
||||
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()), "p", "stop"
|
||||
if "iterative research process" in system:
|
||||
return next(decisions), "d", "stop"
|
||||
captured.update(kwargs)
|
||||
research_db.set_report_progress(run["id"], "# Report\n\nGrounded text.")
|
||||
return "# Report\n\nGrounded text.", "s", "stop"
|
||||
|
||||
def fake_tool(name, arguments, *a, **k):
|
||||
return "page body" if arguments.get("url") else _two_source_search()
|
||||
|
||||
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)))
|
||||
|
||||
assert captured.get("phase") == "synthesis"
|
||||
assert "repetition_penalty" not in captured
|
||||
|
||||
|
||||
def test_auto_scrape_respects_char_budgets(research_home, monkeypatch):
|
||||
worker, supervisor = _bare_supervisor(monkeypatch)
|
||||
_patch_web_rank(monkeypatch)
|
||||
# space-separated so page cleaning keeps it (a single 50k-char token is stripped as junk)
|
||||
monkeypatch.setattr(worker, "execute_tool", lambda *a, **k: "yy " * 20_000)
|
||||
step_sources = [
|
||||
{"url": f"https://s{i}.example.com", "title": f"S{i}"} for i in range(3)
|
||||
]
|
||||
section, fetched = asyncio.run(
|
||||
supervisor._auto_scrape_sources(
|
||||
{"id": "run-x"}, "question", step_sources, set(),
|
||||
tool_timeout = 10, website_policy = None,
|
||||
)
|
||||
)
|
||||
# the folded evidence is bounded chunks, not the 150k of raw page bodies
|
||||
# (the retrieved chunk section is capped at _AUTO_SCRAPE_TOTAL_CHARS; a short fixed
|
||||
# header is prepended on top)
|
||||
assert "<chunk" in section
|
||||
assert len(section) <= worker._AUTO_SCRAPE_TOTAL_CHARS + 200
|
||||
assert len(fetched) == worker._AUTO_SCRAPE_TOP_K
|
||||
notes = [f"### Step\nInput: q\nResult:\n{section[:12_000]}"]
|
||||
assert len(worker._bounded_synthesis_evidence(notes)) <= worker._MAX_SYNTHESIS_EVIDENCE_CHARS
|
||||
|
||||
|
||||
def test_auto_scrape_falls_back_when_no_relevant_chunks(research_home, monkeypatch):
|
||||
# When hybrid retrieval surfaces nothing above the floor (covered in test_web_rank.py),
|
||||
# the step yields no scraped section and the caller keeps the snippet evidence.
|
||||
worker, supervisor = _bare_supervisor(monkeypatch)
|
||||
_patch_web_rank(monkeypatch, retrieve = lambda *a, **k: ("", []))
|
||||
monkeypatch.setattr(worker, "execute_tool", lambda *a, **k: "unrelated boilerplate content")
|
||||
step_sources = [{"url": "https://s.example.com", "title": "S"}]
|
||||
section, fetched = asyncio.run(
|
||||
supervisor._auto_scrape_sources(
|
||||
{"id": "run-x"}, "find the special token", step_sources, set(),
|
||||
tool_timeout = 10, website_policy = None,
|
||||
)
|
||||
)
|
||||
assert section == ""
|
||||
assert fetched == []
|
||||
|
||||
|
||||
def test_clean_scraped_text_strips_nav_and_encoded_links():
|
||||
from core import research_runs as worker
|
||||
|
||||
raw = (
|
||||
"# Qwen\n"
|
||||
"* [العربية](https://ar.wikipedia.org/wiki/%D9%83%D9%88%D9%8A%D9%86_%D9%86%D9%85)\n"
|
||||
"* [Deutsch](https://de.wikipedia.org/wiki/Qwen)\n"
|
||||
"[Qwen](/Qwen) 's Collections\n"
|
||||
"[Qwen-AgentWorld](/collections/Qwen/qwen-agentworld)\n"
|
||||
"BaseModelAndInstructionTuning.html?q=base%2Cmodels&sa=D&sntz=1&usg=AOvVaw2JZPpIYwRrXNjGnFtOuS-H\n"
|
||||
"Qwen2.5 is released under the [Apache 2.0](https://apache.org/licenses) license, "
|
||||
"which permits commercial use and redistribution.\n"
|
||||
"The maximum context length is 131072 tokens.\n"
|
||||
)
|
||||
cleaned = worker._clean_scraped_text(raw)
|
||||
|
||||
# nav sidebars, encoded-URL lists, bare link menus, and tracking-URL tokens are gone
|
||||
assert "العربية" not in cleaned
|
||||
assert "ar.wikipedia" not in cleaned
|
||||
assert "AgentWorld" not in cleaned
|
||||
assert "'s Collections" not in cleaned
|
||||
assert "AOvVaw2" not in cleaned
|
||||
# real prose with an inline link survives
|
||||
assert "Apache 2.0" in cleaned
|
||||
assert "131072 tokens" in cleaned
|
||||
|
||||
|
||||
def test_auto_scrape_skips_already_fetched_urls(research_home, monkeypatch):
|
||||
worker, supervisor = _bare_supervisor(monkeypatch)
|
||||
_patch_web_rank(monkeypatch)
|
||||
called = []
|
||||
|
||||
def fake_tool(name, arguments, *args, **kwargs):
|
||||
called.append(arguments["url"])
|
||||
return "body for " + arguments["url"]
|
||||
|
||||
monkeypatch.setattr(worker, "execute_tool", fake_tool)
|
||||
step_sources = [
|
||||
{"url": "https://x.example.com", "title": "X"},
|
||||
{"url": "https://y.example.com", "title": "Y"},
|
||||
]
|
||||
section, fetched = asyncio.run(
|
||||
supervisor._auto_scrape_sources(
|
||||
{"id": "run-x"},
|
||||
"question",
|
||||
step_sources,
|
||||
{"https://x.example.com"},
|
||||
tool_timeout = 10,
|
||||
website_policy = None,
|
||||
)
|
||||
)
|
||||
assert called == ["https://y.example.com"]
|
||||
assert fetched == ["https://y.example.com"]
|
||||
assert "https://x.example.com" not in section
|
||||
|
||||
|
||||
def test_recovered_running_research_resumes_durable_progress(research_home, monkeypatch):
|
||||
from core import research_runs as worker
|
||||
|
||||
|
|
|
|||
116
studio/backend/tests/test_web_rank.py
Normal file
116
studio/backend/tests/test_web_rank.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
"""Unit tests for the ephemeral web-RAG used by deep research auto-read.
|
||||
|
||||
These run the *real* Studio RAG store + hybrid retrieval + formatter against a temporary
|
||||
rag.db (so the ingest -> 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 "<chunk" in rendered
|
||||
assert "LoRA" in rendered
|
||||
assert sources and sources[0]["citationId"] == 1
|
||||
# source attribution is the page title, via Studio's formatter
|
||||
assert 'source="LoRA"' in rendered
|
||||
|
||||
|
||||
def test_min_score_floor_drops_irrelevant(rag_home):
|
||||
pages = [
|
||||
{"text": "LoRA adapters reduce trainable parameters for fine tuning.", "url": "https://a"},
|
||||
{"text": "Completely separate cooking recipe with onions and garlic.", "url": "https://b"},
|
||||
]
|
||||
rendered, _ = web_rank.retrieve_web_chunks(pages, "lora fine tuning", top_n = 5, min_score = 0.5)
|
||||
assert "cooking" not in rendered.lower()
|
||||
assert "lora" in rendered.lower()
|
||||
|
||||
|
||||
def test_char_budget_caps_kept_chunks(rag_home):
|
||||
# ~2000 words -> 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("<chunk id") >= 2
|
||||
assert 1 <= capped.count("<chunk id") < full.count("<chunk id")
|
||||
|
||||
|
||||
def test_empty_and_invalid_inputs_return_empty(rag_home):
|
||||
assert web_rank.retrieve_web_chunks([], "lora", top_n = 5, min_score = 0.1) == ("", [])
|
||||
assert web_rank.retrieve_web_chunks([{"text": " "}], "lora", top_n = 5, min_score = 0.1) == ("", [])
|
||||
assert web_rank.retrieve_web_chunks([{"text": "lora"}], "", top_n = 5, min_score = 0.1) == ("", [])
|
||||
assert web_rank.retrieve_web_chunks([{"text": "lora"}], "lora", top_n = 0, min_score = 0.1) == ("", [])
|
||||
|
||||
|
||||
def test_ephemeral_scope_is_cleaned_up(rag_home):
|
||||
pages = [{"text": "LoRA low-rank adaptation fine tuning.", "title": "LoRA", "url": "https://a"}]
|
||||
rendered, _ = web_rank.retrieve_web_chunks(pages, "lora", top_n = 5, min_score = 0.0)
|
||||
assert "<chunk" in rendered
|
||||
# nothing from the auto-read is left in the store
|
||||
assert _scope_rows(rag_home) == (0, 0)
|
||||
Loading…
Add table
Add a link
Reference in a new issue