From 2b75062d380cfefbe0a10d0fd8cce9ed6e89bf13 Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Thu, 16 Jul 2026 15:48:17 +0200 Subject: [PATCH] fix: tick quotes --- .../components/engines/citations/utils.py | 177 +++++++++-------- .../engines/test_citation_chunk_boundaries.py | 118 +++++++++++ .../test_citation_parser_robustness.py | 155 +++++++++++++++ .../engines/test_citation_streaming_quotes.py | 43 ++++ ...rocessor_citations_streaming_robustness.py | 185 ++++++++++++++++++ 5 files changed, 599 insertions(+), 79 deletions(-) create mode 100644 tests/engines/test_citation_chunk_boundaries.py create mode 100644 tests/engines/test_citation_parser_robustness.py create mode 100644 tests/engines/test_citation_streaming_quotes.py create mode 100644 tests/engines/test_processor_citations_streaming_robustness.py diff --git a/private_gpt/components/engines/citations/utils.py b/private_gpt/components/engines/citations/utils.py index 0e1f16dc..48c5ddb6 100644 --- a/private_gpt/components/engines/citations/utils.py +++ b/private_gpt/components/engines/citations/utils.py @@ -242,6 +242,8 @@ def _extract_citations_from_text( return cites + + def extract_citations_by_original_text( text: str, documents: list[Document], @@ -256,63 +258,96 @@ def extract_citations_by_original_text( citation_indices = citation_indices or {} result = "" - buffer = "" start_len = len(start_token) end_len = len(end_token) # Model can generate brackets not normalized 【 text = text.replace("【", start_token).replace("】", end_token) - # Iterate through the text to remove malformed citations and save correct citations + # Iterate through the text to remove malformed citations and save correct citations. + # Backticks directly wrapping a citation are formatting noise and are not emitted. i = 0 docs = [] + citation_placeholders: list[str] = [] + code_delimiter: str | None = None + citation_wrapper_delimiter: str | None = None + while i < len(text): - # If we have a buffer, try to process it first - if buffer and start_token in buffer and end_token in buffer: - citation_start = buffer.find(start_token) - citation_end = buffer.find(end_token) + end_len - - citation_text = buffer[citation_start:citation_end] - node_ids = [ - id.strip() - for id in citation_text[len(start_token) : -len(end_token)].split( - split_token - ) - ] - - valid_docs = [] - for node_id in node_ids: - doc = next( - (doc for doc in documents if doc.id.lower() == node_id.lower()), - None, - ) - if doc: - valid_docs.append(doc) - - if valid_docs: - docs.extend(valid_docs) - result += split_token.join( - f"{start_token}{doc.id}{end_token}" for doc in valid_docs - ) - # Clear buffer and continue from after citation - rest_of_buffer = buffer[citation_end:] - if rest_of_buffer and rest_of_buffer[0] == "`": - buffer = rest_of_buffer[1:] - else: - buffer = rest_of_buffer - continue - else: - # No valid docs in citation, treat as regular text - result += buffer - buffer = "" - continue - if text[i] == "`": - if buffer: - buffer += text[i] - else: - buffer = text[i] - i += 1 + delimiter_end = i + 1 + while delimiter_end < len(text) and text[delimiter_end] == "`": + delimiter_end += 1 + delimiter = text[i:delimiter_end] + + if citation_wrapper_delimiter == delimiter: + citation_wrapper_delimiter = None + i = delimiter_end + continue + + if code_delimiter == delimiter: + result += delimiter + code_delimiter = None + i = delimiter_end + continue + + if delimiter_end == len(text): + if is_final: + result += delimiter + break + + if text[delimiter_end : delimiter_end + start_len] == start_token: + citation_start = delimiter_end + citation_end = citation_start + start_len + while ( + citation_end < len(text) + and text[citation_end : citation_end + end_len] != end_token + ): + citation_end += 1 + + if citation_end >= len(text): + break + + node_ids = [ + node_id.strip() + for node_id in text[ + citation_start + start_len : citation_end + ].split(split_token) + ] + valid_docs = [ + doc + for node_id in node_ids + if ( + doc := next( + ( + document + for document in documents + if document.id.lower() == node_id.lower() + ), + None, + ) + ) + ] + if valid_docs: + placeholders = [] + for doc in valid_docs: + placeholder_index = "".join( + f"n{digit}" for digit in str(len(docs)) + ) + placeholder = f"\ue000citation{placeholder_index}\ue001" + docs.append(doc) + citation_placeholders.append(placeholder) + placeholders.append(placeholder) + result += split_token.join(placeholders) + i = citation_end + end_len + if text[i : i + len(delimiter)] == delimiter: + i += len(delimiter) + else: + citation_wrapper_delimiter = delimiter + continue + + result += delimiter + code_delimiter = delimiter + i = delimiter_end elif text[i : i + start_len] == start_token: # Check if we have a complete citation j = i + start_len @@ -332,22 +367,18 @@ def extract_citations_by_original_text( if doc: valid_docs.append(doc) if valid_docs: - if buffer and buffer != "`": - result += buffer - buffer = "" - docs.extend(valid_docs) - result += split_token.join( - f"{start_token}{doc.id}{end_token}" for doc in valid_docs - ) - next_pos = j + end_len - if next_pos < len(text) and text[next_pos] == "`": - i = next_pos + 1 - else: - i = j + end_len + placeholders = [] + for doc in valid_docs: + placeholder_index = "".join( + f"n{digit}" for digit in str(len(docs)) + ) + placeholder = f"\ue000citation{placeholder_index}\ue001" + docs.append(doc) + citation_placeholders.append(placeholder) + placeholders.append(placeholder) + result += split_token.join(placeholders) + i = j + end_len else: - if buffer and buffer != "`": - result += buffer - buffer = "" # No valid docs in citation, treat as regular text result += text[i : j + end_len] i = j + end_len @@ -359,19 +390,9 @@ def extract_citations_by_original_text( i = len(text) continue else: - if not buffer: - result += text[i] - else: - buffer += text[i] + result += text[i] i += 1 - # Don't output buffer if: - # 1. It's just a backtick (could be start of citation) and we are not final - # 2. It contains start token but not end token (incomplete citation) - if buffer and start_token not in buffer: - if is_final or buffer != "`": - result += buffer - # Process citations pattern = re.compile( rf"{re.escape(start_token)}?[A-Z0-9]{shorter_id_length}{re.escape(end_token)}?" @@ -395,7 +416,9 @@ def extract_citations_by_original_text( max_index = max(citation_indices.values(), default=-1) current_index = max_index + 1 processed_docs = [] - for i, doc in enumerate(docs): + for i, (doc, placeholder) in enumerate( + zip(docs, citation_placeholders, strict=True) + ): if doc.id_ not in processed_docs: processed_docs.append(doc.id_) @@ -409,11 +432,7 @@ def extract_citations_by_original_text( citation_indices[doc.id_] = index citation = format_cite(i, doc, index) - result = result.replace( - f"{start_token}{doc.id}{end_token}", - citation, - 1, - ) + result = result.replace(placeholder, citation, 1) return result, _extract_citations_from_text(result), citation_indices diff --git a/tests/engines/test_citation_chunk_boundaries.py b/tests/engines/test_citation_chunk_boundaries.py new file mode 100644 index 00000000..7d4df9fd --- /dev/null +++ b/tests/engines/test_citation_chunk_boundaries.py @@ -0,0 +1,118 @@ +import random + +import pytest + +from private_gpt.components.engines.citations.types import Document +from private_gpt.components.engines.citations.utils import ( + extract_citations_by_original_text, +) + + +TEXT = ( + "Format: `[XXXX]`. Correct: `[AB12]`. " + "Invalid: `[[AB12]]`, `(AB12)`. " + "Consolidate: `[AB12], [CD34]`. Final [EF56]." +) + + +def create_documents() -> list[Document]: + return [ + Document( + type="document", + id_=f"source-{citation_id}", + shorter_id=citation_id, + document_id=f"artifact-{citation_id}", + text=citation_id, + ) + for citation_id in ("AB12", "CD34", "EF56") + ] + + +def stream_chunks(chunks: list[str]) -> tuple[str, int]: + documents = create_documents() + current_text = "" + sent_text = "" + emitted_text = "" + citation_indices: dict[str, int] = {} + citation_count = 0 + + for chunk in chunks: + current_text += chunk + cleaned_text, citations, citation_indices = extract_citations_by_original_text( + current_text, + documents, + citation_indices=citation_indices, + ) + assert cleaned_text.startswith(sent_text) + emitted_text += cleaned_text[len(sent_text) :] + sent_text = cleaned_text + citation_count = len(citations) + + final_text, citations, _ = extract_citations_by_original_text( + current_text, + documents, + citation_indices=citation_indices, + is_final=True, + ) + assert final_text.startswith(sent_text) + emitted_text += final_text[len(sent_text) :] + return emitted_text, len(citations) if citations else citation_count + + +def expected_result() -> tuple[str, int]: + formatted, citations, _ = extract_citations_by_original_text( + TEXT, + create_documents(), + is_final=True, + ) + return formatted, len(citations) + + +def test_every_two_chunk_split_matches_one_shot_result() -> None: + expected_text, expected_citations = expected_result() + + for split_at in range(len(TEXT) + 1): + streamed_text, citation_count = stream_chunks( + [TEXT[:split_at], TEXT[split_at:]] + ) + assert streamed_text == expected_text, f"split_at={split_at}" + assert citation_count == expected_citations, f"split_at={split_at}" + + +def test_character_by_character_stream_matches_one_shot_result() -> None: + assert stream_chunks(list(TEXT)) == expected_result() + + +@pytest.mark.parametrize("seed", range(20)) +def test_random_chunking_matches_one_shot_result(seed: int) -> None: + random_generator = random.Random(seed) + chunks = [] + position = 0 + while position < len(TEXT): + chunk_size = random_generator.randint(1, 12) + chunks.append(TEXT[position : position + chunk_size]) + position += chunk_size + + assert stream_chunks(chunks) == expected_result() + + +@pytest.mark.parametrize( + "partial", + [ + "Prefix [", + "Prefix [A", + "Prefix [AB", + "Prefix [AB1", + "Prefix `[AB12", + "Prefix ``[AB12", + "Prefix ```[AB12", + ], +) +def test_partial_citation_never_leaks_during_stream(partial: str) -> None: + formatted, citations, _ = extract_citations_by_original_text( + partial, + create_documents(), + ) + + assert formatted == "Prefix " + assert citations == [] diff --git a/tests/engines/test_citation_parser_robustness.py b/tests/engines/test_citation_parser_robustness.py new file mode 100644 index 00000000..76ccfbd8 --- /dev/null +++ b/tests/engines/test_citation_parser_robustness.py @@ -0,0 +1,155 @@ +import pytest + +from private_gpt.components.engines.citations.types import Document +from private_gpt.components.engines.citations.utils import ( + extract_citations_by_original_text, + format_cite, +) + + +def create_document(citation_id: str) -> Document: + return Document( + type="document", + id_=f"source-{citation_id}", + shorter_id=citation_id, + document_id=f"artifact-{citation_id}", + text=f"Content for {citation_id}", + ) + + +def test_repeated_citation_keeps_index_but_emits_each_occurrence() -> None: + document = create_document("AB12") + + formatted, citations, indices = extract_citations_by_original_text( + "First [AB12], repeated [AB12].", + [document], + ) + + assert formatted == ( + f"First {format_cite(0, document, 0)}, repeated {format_cite(1, document, 0)}." + ) + assert [citation.value["index"] for citation in citations] == ["0", "0"] + assert indices == {document.id_: 0} + + +def test_mixed_known_and_unknown_consolidated_citation_keeps_known_only() -> None: + document = create_document("AB12") + + formatted, citations, _ = extract_citations_by_original_text( + "Claim [UNKNOWN, AB12, MISSING].", + [document], + ) + + assert formatted == f"Claim {format_cite(0, document, 0)}." + assert len(citations) == 1 + + +def test_citation_lookup_is_case_insensitive() -> None: + document = create_document("AB12") + + formatted, citations, _ = extract_citations_by_original_text( + "Claim [ab12].", + [document], + ) + + assert formatted == f"Claim {format_cite(0, document, 0)}." + assert len(citations) == 1 + + +def test_unicode_citation_brackets_are_normalized() -> None: + document = create_document("AB12") + + formatted, citations, _ = extract_citations_by_original_text( + "Claim 【AB12】.", + [document], + ) + + assert formatted == f"Claim {format_cite(0, document, 0)}." + assert len(citations) == 1 + + +@pytest.mark.parametrize("delimiter", ["`", "``", "```"]) +def test_backtick_wrapped_citation_removes_matching_delimiter(delimiter: str) -> None: + document = create_document("AB12") + + formatted, citations, _ = extract_citations_by_original_text( + f"Claim {delimiter}[AB12]{delimiter}.", + [document], + ) + + assert formatted == f"Claim {format_cite(0, document, 0)}." + assert len(citations) == 1 + + +@pytest.mark.parametrize( + "garbage", + [ + "[]", + "[ ]", + "[[AB12]]", + "(AB12)", + "[AB-12]", + "[TOO-LONG]", + "[UNKNOWN]", + "prefix AB12 suffix", + ], +) +def test_non_citation_garbage_is_preserved(garbage: str) -> None: + document = create_document("AB12") + + formatted, citations, _ = extract_citations_by_original_text( + garbage, + [document], + ) + + assert formatted == garbage + assert citations == [] + + +def test_incomplete_citation_is_withheld_with_no_false_citation() -> None: + document = create_document("AB12") + + formatted, citations, _ = extract_citations_by_original_text( + "Safe prefix [AB1", + [document], + ) + + assert formatted == "Safe prefix " + assert citations == [] + + +@pytest.mark.xfail( + strict=True, + reason="Model output can currently collide with the internal citation placeholder.", +) +def test_placeholder_like_model_output_does_not_capture_real_citation() -> None: + document = create_document("AB12") + model_text = "Literal \ue000citationn0\ue001 then [AB12]." + + formatted, citations, _ = extract_citations_by_original_text( + model_text, + [document], + ) + + assert formatted == ( + "Literal \ue000citationn0\ue001 then " + f"{format_cite(0, document, 0)}." + ) + assert len(citations) == 1 + + +def test_existing_indices_continue_without_renumbering() -> None: + first = create_document("AB12") + second = create_document("CD34") + + formatted, citations, indices = extract_citations_by_original_text( + "Existing [AB12], new [CD34].", + [first, second], + citation_indices={first.id_: 7}, + ) + + assert formatted == ( + f"Existing {format_cite(0, first, 7)}, new {format_cite(1, second, 8)}." + ) + assert [citation.value["index"] for citation in citations] == ["7", "8"] + assert indices == {first.id_: 7, second.id_: 8} diff --git a/tests/engines/test_citation_streaming_quotes.py b/tests/engines/test_citation_streaming_quotes.py new file mode 100644 index 00000000..be9e2df5 --- /dev/null +++ b/tests/engines/test_citation_streaming_quotes.py @@ -0,0 +1,43 @@ +from private_gpt.components.engines.citations.types import Document +from private_gpt.components.engines.citations.utils import ( + extract_citations_by_original_text, +) + + +def test_streaming_citation_examples_do_not_rewrite_previous_text() -> None: + text = ( + "Format: `[XXXX]`. Correct: `[LLFB]`. " + "Invalid: `[[LLFB]]`, `(LLFB)`. " + "Consolidate: `[LLFB], [CPCY]`. End." + ) + documents = [ + Document( + type="document", + id_=f"source-{citation_id}", + shorter_id=citation_id, + document_id="artifact", + text=citation_id, + ) + for citation_id in ("LLFB", "CPCY") + ] + current_text = "" + sent_text = "" + emitted_text = "" + citation_indices: dict[str, int] = {} + + for character in text: + current_text += character + cleaned_text, _, citation_indices = extract_citations_by_original_text( + current_text, + documents, + citation_indices=citation_indices, + ) + assert cleaned_text.startswith(sent_text) + emitted_text += cleaned_text[len(sent_text) :] + sent_text = cleaned_text + + assert "Format: `[XXXX]`." in emitted_text + assert "Invalid: `[[LLFB]]`, `(LLFB)`." in emitted_text + assert "Consolidate: `. End." not in emitted_text diff --git a/tests/engines/test_processor_citations_streaming_robustness.py b/tests/engines/test_processor_citations_streaming_robustness.py new file mode 100644 index 00000000..245ca826 --- /dev/null +++ b/tests/engines/test_processor_citations_streaming_robustness.py @@ -0,0 +1,185 @@ +import asyncio +from collections.abc import AsyncGenerator + +import pytest + +from private_gpt.components.chat.processors.events.citations.citations import ( + process_citations, +) +from private_gpt.components.engines.citations.types import Document +from private_gpt.components.engines.citations.utils import format_cite +from private_gpt.events.models import ( + Event, + RawContentBlockDeltaEvent, + RawContentBlockStartEvent, + RawContentBlockStopEvent, + TextBlock, + TextDelta, + ThinkingBlock, + ThinkingDelta, +) + + +def create_document(citation_id: str) -> Document: + return Document( + type="document", + id_=f"source-{citation_id}", + shorter_id=citation_id, + document_id=f"artifact-{citation_id}", + text=citation_id, + ) + + +async def delayed_events( + events: list[Event], delay_seconds: float = 0 +) -> AsyncGenerator[Event, None]: + for event in events: + await asyncio.sleep(delay_seconds) + yield event + + +async def collect_text_stream( + chunks: list[str], + documents: list[Document], + delay_seconds: float = 0, +) -> tuple[str, int]: + events: list[Event] = [ + RawContentBlockStartEvent(block_id="text", content_block=TextBlock(text="")), + *[ + RawContentBlockDeltaEvent( + block_id="text", + delta=TextDelta(text=chunk), + ) + for chunk in chunks + ], + RawContentBlockStopEvent(block_id="text"), + ] + output = "" + citation_count = 0 + + async for event in process_citations( + delayed_events(events, delay_seconds), + lambda **kwargs: documents, + ): + if not isinstance(event, RawContentBlockDeltaEvent): + continue + if not isinstance(event.delta, TextDelta): + continue + output += event.delta.text or "" + citation_count += len(event.delta.citations or []) + + return output, citation_count + + +async def collect_thinking_stream( + chunks: list[str], documents: list[Document] +) -> tuple[str, int]: + events: list[Event] = [ + RawContentBlockStartEvent( + block_id="thinking", + content_block=ThinkingBlock(thinking="", signature=""), + ), + *[ + RawContentBlockDeltaEvent( + block_id="thinking", + delta=ThinkingDelta(thinking=chunk), + ) + for chunk in chunks + ], + RawContentBlockStopEvent(block_id="thinking"), + ] + output = "" + citation_count = 0 + + async for event in process_citations( + delayed_events(events), + lambda **kwargs: documents, + ): + if not isinstance(event, RawContentBlockDeltaEvent): + continue + if not isinstance(event.delta, ThinkingDelta): + continue + output += event.delta.thinking or "" + citation_count += len(event.delta.citations or []) + + return output, citation_count + + +@pytest.mark.asyncio +async def test_every_event_split_inside_wrapped_citation_is_stable() -> None: + document = create_document("AB12") + text = "Before `[AB12]` after." + expected = f"Before {format_cite(0, document, 0)} after." + + for split_at in range(len(text) + 1): + output, citation_count = await collect_text_stream( + [text[:split_at], text[split_at:]], + [document], + ) + assert output == expected, f"split_at={split_at}" + assert citation_count == 1, f"split_at={split_at}" + + +@pytest.mark.asyncio +async def test_delayed_character_stream_preserves_output_and_citations() -> None: + first = create_document("AB12") + second = create_document("CD34") + text = "Before `[AB12], [CD34]` after." + + output, citation_count = await collect_text_stream( + list(text), + [first, second], + delay_seconds=0.0001, + ) + + assert output == ( + f"Before {format_cite(0, first, 0)}, {format_cite(1, second, 1)} after." + ) + assert citation_count == 2 + + +@pytest.mark.asyncio +async def test_thinking_stream_uses_same_citation_semantics() -> None: + document = create_document("AB12") + + output, citation_count = await collect_thinking_stream( + ["Reasoning `", "[AB", "12]", "` complete."], + [document], + ) + + assert output == f"Reasoning {format_cite(0, document, 0)} complete." + assert citation_count == 1 + + +@pytest.mark.asyncio +async def test_llm_garbage_around_valid_citation_is_not_lost() -> None: + document = create_document("AB12") + chunks = [ + "Trash `[[AB12]]`, `(AB12)`, [UNKNOWN]. ", + "Valid `[AB", + "12]` end.", + ] + + output, citation_count = await collect_text_stream(chunks, [document]) + + assert output == ( + "Trash `[[AB12]]`, `(AB12)`, [UNKNOWN]. " + f"Valid {format_cite(0, document, 0)} end." + ) + assert citation_count == 1 + + +@pytest.mark.asyncio +async def test_repeated_citation_across_delayed_events_reuses_index() -> None: + document = create_document("AB12") + + output, citation_count = await collect_text_stream( + ["First [AB12]", ", second ", "[AB12]."], + [document], + delay_seconds=0.0001, + ) + + assert output == ( + f"First {format_cite(0, document, 0)}, second {format_cite(1, document, 0)}." + ) + assert citation_count == 2