3c658e4666
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.
116 lines
4.6 KiB
Python
116 lines
4.6 KiB
Python
"""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)
|