diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index c229720e84..4d677c07b7 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -42,7 +42,10 @@ _SOURCES_HEADING = re.compile( _NUMBERED_CITATION = re.compile(r"(?\s]+)>") _RAW_URL = re.compile(r"https?://[^\s<>]+") -_DOCUMENT_CITATION = re.compile(r"\[Document:(?:[^\[\]]+|\[[^\[\]]*\])*\]") +# Unrolled rather than the equivalent (?:[^\[\]]+|\[[^\[\]]*\])* : that alternation backtracks +# catastrophically on an unterminated "[Document:" (ordinary malformed model output), and this +# runs on the event loop, so one bad report would stall all of Studio. +_DOCUMENT_CITATION = re.compile(r"\[Document:[^\[\]]*(?:\[[^\[\]]*\][^\[\]]*)*\]") # Wrapper delimiters used in the decision/synthesis prompts. Any occurrence inside # untrusted evidence is escaped so gathered content cannot close a block early. _PROMPT_DELIMITER_TAGS = re.compile( @@ -645,6 +648,18 @@ def _split_rag_result(result: str) -> tuple[str, list[dict[str, Any]]]: return text.rstrip(), sources +def _citation_title(source: dict, fallback: str) -> str: + """Title as it may appear in a markdown link label. + + Brackets are stripped because the report prompt tells the model to copy titles verbatim + from the source catalog, and search titles routinely carry one ("[PDF] Annual Report"). + A bracket inside the label makes the citation unmatchable, so the catalog and the + citation writer must agree on the same stripped form. + """ + title = str(source.get("title") or fallback).replace("[", "").replace("]", "").strip() + return title or fallback + + def _trim_url_tail(raw: str) -> str: """Strip trailing prose punctuation that ``_RAW_URL`` swallowed. @@ -688,9 +703,9 @@ def _validate_report_sources(report: str, sources: list[dict]) -> str: source = source_by_url.get(url) if source is None: return None - title = str(source.get("title") or url).replace("[", "").replace("]", "").strip() + title = _citation_title(source, url) token = f"\x00research-citation-{len(placeholders)}\x00" - placeholders[token] = f"[{title or url}]({_escape_link_destination(url)})" + placeholders[token] = f"[{title}]({_escape_link_destination(url)})" return token def replace_markdown_links(text: str) -> str: @@ -1627,30 +1642,33 @@ class ResearchSupervisor: ) for source in document_sources } + # Mirrors the live loop: evidence must hold only chunks that made it into the + # catalog, else the validator strips citations to the rest and synthesis is left + # building claims on uncataloged document text. + accepted_rag_sources = [] for source in restored_rag_sources: source_key = str( source.get("chunkId") or f"{source.get('documentId') or source.get('filename')}:{source.get('page') or ''}" ) - if ( - source_key in document_source_keys - or len(sources) + len(document_sources) >= max_sources - ): - continue - written = await asyncio.to_thread( - db.upsert_document_source, - run["id"], - int(step["position"]), - source, - self.worker_id, - ) - await self._check_worker_write(run["id"], written) - document_source_keys.add(source_key) - document_sources.append({**source, "stepPosition": step["position"]}) + if source_key not in document_source_keys: + if len(sources) + len(document_sources) >= max_sources: + continue + written = await asyncio.to_thread( + db.upsert_document_source, + run["id"], + int(step["position"]), + source, + self.worker_id, + ) + await self._check_worker_write(run["id"], written) + document_source_keys.add(source_key) + document_sources.append({**source, "stepPosition": step["position"]}) + accepted_rag_sources.append(source) rag_evidence = "\n".join( f"{item.get('filename') or 'Document'}: " f"{item.get('text') or item.get('snippet') or ''}" - for item in restored_rag_sources + for item in accepted_rag_sources ) title = str(step.get("title") or "Recovered research step") notes.append( @@ -1671,7 +1689,7 @@ class ResearchSupervisor: for position in range(start_position, max_steps): await self._check_active(run["id"]) source_catalog = "\n".join( - f"- {source.get('title') or source['url']} | {source['url']} | " + f"- {_citation_title(source, source['url'])} | {source['url']} | " f"{source.get('snippet') or ''}" for source in sources ) @@ -1837,6 +1855,12 @@ class ResearchSupervisor: f"{source.get('text') or source.get('snippet') or ''}" for source in accepted_rag_sources ) + elif rag_sources: + # Every chunk was refused by the source cap, so none has a catalog entry and the + # validator would strip any citation to it. Drop the evidence rather than let + # synthesis build claims on it. Gated on rag_sources so a text-only KB reply + # ("No documents are attached to this chat.") is still passed through. + rag_result = "" rag_sources = accepted_rag_sources step_sources = [] for match in _URL_BLOCK.finditer(result if action["action"] == "search" else ""): @@ -1943,7 +1967,7 @@ class ResearchSupervisor: await self._check_worker_write(run["id"], seq is not None) await self._check_active(run["id"]) source_catalog = "\n".join( - f"{index}. Title: {source.get('title') or source['url']}\n URL: {source['url']}" + f"{index}. Title: {_citation_title(source, source['url'])}\n URL: {source['url']}" for index, source in enumerate(sources, 1) ) document_source_catalog = "\n".join( diff --git a/studio/backend/tests/test_research_runs_hardening.py b/studio/backend/tests/test_research_runs_hardening.py index afff5b6f00..40316f683b 100644 --- a/studio/backend/tests/test_research_runs_hardening.py +++ b/studio/backend/tests/test_research_runs_hardening.py @@ -6,6 +6,7 @@ import pytest from core.research_runs import ( + _citation_title, _escape_link_destination, _sanitize_public_query, _shield_untrusted, @@ -121,6 +122,32 @@ def test_document_citation_strips_unknown_source_with_brackets(): assert out == "Ghost cite end." +def test_document_citation_regex_does_not_backtrack_catastrophically(): + # An unterminated "[Document:" with no later bare "]" is ordinary malformed model output, + # which is exactly what this sanitizer exists to handle. The old alternation took longer + # than the age of the universe on one line, and it runs on the event loop. + import time + + report = "Revenue rose 12 percent [Document: q3_report.pdf, p. 12 and margins improved." + start = time.perf_counter() + _validate_report_document_sources(report, [{"filename": "q3_report.pdf", "page": 12}]) + assert time.perf_counter() - start < 1.0 + # And a long tail stays linear rather than exponential. + start = time.perf_counter() + _validate_report_document_sources("[Document: " + "a" * 20_000, []) + assert time.perf_counter() - start < 1.0 + + +def test_citation_title_strips_brackets_for_catalog_and_citation(): + # Search titles routinely carry a bracketed prefix ("[PDF] ..."), and the report prompt tells + # the model to copy the catalog title verbatim into the link label, where a bracket makes the + # citation unmatchable. The catalog and the citation writer share this helper so they cannot + # offer a label the validator then fails to match. + assert _citation_title({"title": "[PDF] Annual Report 2024"}, "https://x/a") == "PDF Annual Report 2024" + assert _citation_title({"title": "[]"}, "https://x/a") == "https://x/a" + assert _citation_title({}, "https://x/a") == "https://x/a" + + def _make_payload(**overrides) -> CreateResearchRun: payload = {"threadId": "t1", "userMessageId": "u1", "inferenceRequest": {"model": "m"}} payload.update(overrides) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 5f67c39f3d..c87ef586b8 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -1923,6 +1923,85 @@ def test_recovered_running_research_resumes_durable_progress(research_home, monk assert completed["report"].startswith("# Resumed report") +def test_knowledge_base_evidence_beyond_the_source_cap_is_not_synthesized( + research_home, monkeypatch +): + """A knowledge-base hit that the source cap refuses to persist must not reach synthesis: + it has no document_source_catalog entry, so any citation of it is stripped from the + finished report and the claim it supports would be left unattributed.""" + from core import research_runs as worker + + _create( + rag_scope = {"kb_id": "kb-1", "default_top_k": 4}, + budgets = { + "maxSteps": 3, + "maxSources": 1, + "modelTimeoutSeconds": 30, + "toolTimeoutSeconds": 10, + }, + ) + supervisor = worker.ResearchSupervisor(SimpleNamespace(state = SimpleNamespace(server_port = 1))) + decisions = iter( + ( + json.dumps({"action": "search", "title": "First", "query": "first query"}), + json.dumps({"action": "search", "title": "Second", "query": "second query"}), + json.dumps({"action": "finish", "title": "Enough evidence"}), + ) + ) + synthesis_prompts = [] + report = "# Report\n\nA finding [Document: kept.pdf, p. 1]." + + async def fake_stream_completion(run, messages, **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" + + labels = iter(("kept", "capped")) + + def fake_tool(name, arguments, *args, **kwargs): + if name == "search_knowledge_base": + label = next(labels) + return ( + f"UNCATALOGED_{label.upper()}_KB_TEXT" + + worker.RAG_SOURCES_SENTINEL + + json.dumps( + [ + { + "chunkId": f"doc-{label}:0", + "documentId": f"doc-{label}", + "filename": f"{label}.pdf", + "page": 1, + "text": f"{label} chunk body", + } + ] + ) + ) + return "Title: Alpha\nURL: https://a.example.com\nSnippet: alpha snippet." + + 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))) + + completed = research_db.get_run("run-1") + assert completed["status"] == "completed" + # The cap admitted the first chunk only, so only it may appear in the evidence. + assert [source["filename"] for source in completed["documentSources"]] == ["kept.pdf"] + assert synthesis_prompts, "synthesis must have run" + assert "kept chunk body" in synthesis_prompts[0] + assert "UNCATALOGED_KEPT_KB_TEXT" not in synthesis_prompts[0] + assert "capped chunk body" not in synthesis_prompts[0] + assert "UNCATALOGED_CAPPED_KB_TEXT" not in synthesis_prompts[0] + + def test_create_without_assistant_id_does_not_eagerly_create_message(research_home): from routes.research_runs import CreateResearchRun, create_research_run