diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index a815a8ece4..e8891d137d 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -42,7 +42,7 @@ _SOURCES_HEADING = re.compile( _NUMBERED_CITATION = re.compile(r"(?\s]+)>") _RAW_URL = re.compile(r"https?://[^\s<>]+") -_DOCUMENT_CITATION = re.compile(r"\[Document:[^\]]+\]") +_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( diff --git a/studio/backend/tests/test_research_runs_hardening.py b/studio/backend/tests/test_research_runs_hardening.py index e6a4f687b5..221e1088d0 100644 --- a/studio/backend/tests/test_research_runs_hardening.py +++ b/studio/backend/tests/test_research_runs_hardening.py @@ -111,6 +111,16 @@ def test_document_citation_strips_unknown_source(): assert "not-a-real-file" not in out +def test_document_citation_strips_unknown_source_with_brackets(): + # An invalid citation whose filename contains brackets must be removed whole; the old regex + # stopped at the first ``]`` and left the tail (".pdf, p. 9]") behind. + report = "Ghost cite [Document: invented [final].pdf, p. 9] end." + out = _validate_report_document_sources(report, [{"filename": "real.pdf", "page": 1}]) + assert "invented" not in out + assert ".pdf" not in out + assert out == "Ghost cite end." + + def _make_payload(**overrides) -> CreateResearchRun: payload = {"threadId": "t1", "userMessageId": "u1", "inferenceRequest": {"model": "m"}} payload.update(overrides)