From ce25f73a4818a6471afce1cf27a117a4df8d8514 Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Thu, 16 Jul 2026 16:04:18 +0200 Subject: [PATCH] fix: citations --- .../processors/events/citations/citations.py | 2 +- .../components/engines/citations/parser.py | 97 +++++++++ .../components/engines/citations/utils.py | 193 +----------------- .../components/text_processing/__init__.py | 29 +++ .../components/text_processing/engine.py | 114 +++++++++++ .../components/text_processing/models.py | 59 ++++++ .../components/text_processing/rules.py | 186 +++++++++++++++++ .../components/text_processing/test_engine.py | 131 ++++++++++++ .../engines/test_citation_chunk_boundaries.py | 1 - .../test_citation_parser_robustness.py | 7 +- tests/engines/test_processor_citations.py | 4 +- 11 files changed, 631 insertions(+), 192 deletions(-) create mode 100644 private_gpt/components/engines/citations/parser.py create mode 100644 private_gpt/components/text_processing/__init__.py create mode 100644 private_gpt/components/text_processing/engine.py create mode 100644 private_gpt/components/text_processing/models.py create mode 100644 private_gpt/components/text_processing/rules.py create mode 100644 tests/components/text_processing/test_engine.py diff --git a/private_gpt/components/chat/processors/events/citations/citations.py b/private_gpt/components/chat/processors/events/citations/citations.py index de4614c6..dd996743 100644 --- a/private_gpt/components/chat/processors/events/citations/citations.py +++ b/private_gpt/components/chat/processors/events/citations/citations.py @@ -144,7 +144,7 @@ async def process_citations( delta = TextDelta.from_citations(delta_text, delta_citation) yield RawContentBlockDeltaEvent(block_id=event.block_id, delta=delta) - + send_text = cleaned_text send_citations.extend(delta_citation) diff --git a/private_gpt/components/engines/citations/parser.py b/private_gpt/components/engines/citations/parser.py new file mode 100644 index 00000000..f96cdbc2 --- /dev/null +++ b/private_gpt/components/engines/citations/parser.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from collections.abc import Callable + +from private_gpt.components.engines.citations.types import Document +from private_gpt.components.text_processing import ( + BacktickUnwrapRule, + DelimitedReferenceRule, + IncrementalTextProcessor, + ProcessingContext, +) + +CitationFormatter = Callable[[int, Document, int], str] + + +class CitationTextParser: + def __init__( + self, + documents: list[Document], + formatter: CitationFormatter, + *, + start_token: str, + end_token: str, + separator: str, + identifier_length: int, + citation_indices: dict[str, int] | None = None, + ) -> None: + self._documents_by_id = { + document.id.lower(): document for document in documents + } + self._start_token = start_token + self._end_token = end_token + self._formatter = formatter + self._initial_indices = dict(citation_indices or {}) + self._context = ProcessingContext( + state={ + "citation_indices": dict(self._initial_indices), + "citation_next_index": max(self._initial_indices.values(), default=-1) + + 1, + "citation_occurrence": 0, + } + ) + reference_rule = DelimitedReferenceRule( + start_token=start_token, + end_token=end_token, + separator=separator, + resolve=self._resolve, + render=self._render, + ) + self._processor = IncrementalTextProcessor( + [ + BacktickUnwrapRule(reference_rule), + reference_rule, + ] + ) + + def parse(self, text: str, *, final: bool = False) -> tuple[str, dict[str, int]]: + normalized = text.replace("【", self._start_token).replace( + "】", self._end_token + ) + result = self._processor.process( + normalized, + final=final, + context=self._context, + ) + return result.text, dict(self._context.state["citation_indices"]) + + def _resolve( + self, identifiers: list[str], context: ProcessingContext + ) -> list[Document]: + return [ + self._documents_by_id[identifier.lower()] + for identifier in identifiers + if identifier.lower() in self._documents_by_id + ] + + def _render( + self, documents: list[Document], context: ProcessingContext + ) -> tuple[str, tuple[Document, ...]]: + indices: dict[str, int] = context.state["citation_indices"] + next_index: int = context.state["citation_next_index"] + occurrence: int = context.state["citation_occurrence"] + rendered = [] + + for document in documents: + if document.id_ in indices: + index = indices[document.id_] + else: + index = next_index + indices[document.id_] = index + next_index += 1 + rendered.append(self._formatter(occurrence, document, index)) + occurrence += 1 + + context.state["citation_next_index"] = next_index + context.state["citation_occurrence"] = occurrence + return ",".join(rendered), tuple(documents) diff --git a/private_gpt/components/engines/citations/utils.py b/private_gpt/components/engines/citations/utils.py index 48c5ddb6..bf49c369 100644 --- a/private_gpt/components/engines/citations/utils.py +++ b/private_gpt/components/engines/citations/utils.py @@ -12,6 +12,7 @@ from llama_index.core.schema import MetadataMode, NodeWithScore from private_gpt.components.chat.processors.chat_history.memory.utils.splitting import ( get_user_blocks, ) +from private_gpt.components.engines.citations.parser import CitationTextParser from private_gpt.components.engines.citations.types import Citation, Document from private_gpt.components.ingest.metadata_helper import ( MetadataFlags, @@ -242,8 +243,6 @@ def _extract_citations_from_text( return cites - - def extract_citations_by_original_text( text: str, documents: list[Document], @@ -254,187 +253,17 @@ def extract_citations_by_original_text( citation_indices: dict[str, int] | None = None, is_final: bool = False, ) -> tuple[str, list[Citation], dict[str, int]]: - # Initialize an empty string to store the cleaned text - citation_indices = citation_indices or {} - - result = "" - 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. - # 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 text[i] == "`": - 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 - while j < len(text) and text[j : j + end_len] != end_token: - j += 1 - - if j < len(text) and text[j : j + end_len] == end_token: - node_ids = [ - id.strip() for id in text[i + start_len : j].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: - 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: - # No valid docs in citation, treat as regular text - result += text[i : j + end_len] - i = j + end_len - else: - # Incomplete citation detected: drop the - # citation content and stop processing further. - # This change fixes the issue by not - # appending any incomplete citation text. - i = len(text) - continue - else: - result += text[i] - i += 1 - - # Process citations - pattern = re.compile( - rf"{re.escape(start_token)}?[A-Z0-9]{shorter_id_length}{re.escape(end_token)}?" + parser = CitationTextParser( + documents, + format_cite, + start_token=start_token, + end_token=end_token, + separator=split_token, + identifier_length=shorter_id_length, + citation_indices=citation_indices, ) - for match in pattern.finditer(result): - # If citation is well-formed, skip - word = match.group(0) if match.groups() else "" - if not word or (word.startswith(start_token) and word.endswith(end_token)): - continue - - # Try to find the related document, if no document found, skip - doc = next((doc for doc in documents if doc.id in word), None) - if not doc: - continue - - # Remove doc reference - new_token = word.replace(doc.id, "", 1).lstrip().rstrip() - result = result[: match.start()] + new_token + result[match.end() :] - - # Replace citations with sequential numbers, just first occurrence - max_index = max(citation_indices.values(), default=-1) - current_index = max_index + 1 - processed_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_) - - if doc.id_ in citation_indices: - # If we already have this document, use the existing index - index = citation_indices[doc.id_] - else: - # Otherwise, assign a new index - index = current_index - current_index += 1 - citation_indices[doc.id_] = index - - citation = format_cite(i, doc, index) - result = result.replace(placeholder, citation, 1) - - return result, _extract_citations_from_text(result), citation_indices + result, updated_indices = parser.parse(text, final=is_final) + return result, _extract_citations_from_text(result), updated_indices async def deduplicate_documents_in_history( diff --git a/private_gpt/components/text_processing/__init__.py b/private_gpt/components/text_processing/__init__.py new file mode 100644 index 00000000..717f1d2e --- /dev/null +++ b/private_gpt/components/text_processing/__init__.py @@ -0,0 +1,29 @@ +from private_gpt.components.text_processing.engine import IncrementalTextProcessor +from private_gpt.components.text_processing.models import ( + Action, + ProbeResult, + ProbeStatus, + ProcessDelta, + ProcessingContext, + ProcessResult, +) +from private_gpt.components.text_processing.rules import ( + BacktickUnwrapRule, + DelimitedReferenceRule, + LooseReferenceCleanupRule, + StreamRule, +) + +__all__ = [ + "Action", + "BacktickUnwrapRule", + "DelimitedReferenceRule", + "IncrementalTextProcessor", + "LooseReferenceCleanupRule", + "ProbeResult", + "ProbeStatus", + "ProcessDelta", + "ProcessResult", + "ProcessingContext", + "StreamRule", +] diff --git a/private_gpt/components/text_processing/engine.py b/private_gpt/components/text_processing/engine.py new file mode 100644 index 00000000..04d61331 --- /dev/null +++ b/private_gpt/components/text_processing/engine.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from copy import deepcopy +from typing import TYPE_CHECKING + +from private_gpt.components.text_processing.models import ( + Action, + ProbeStatus, + ProcessDelta, + ProcessingContext, + ProcessResult, +) + +if TYPE_CHECKING: + from private_gpt.components.text_processing.rules import StreamRule + + +class IncrementalTextProcessor: + def __init__( + self, + rules: list[StreamRule], + initial_state: dict[str, object] | None = None, + ) -> None: + self._rules = sorted(rules, key=lambda rule: rule.priority, reverse=True) + self._initial_state = deepcopy(initial_state or {}) + self._source = "" + self._emitted = "" + self._metadata_count = 0 + self.context = ProcessingContext(state=deepcopy(self._initial_state)) + + def process( + self, + text: str, + *, + final: bool = False, + context: ProcessingContext | None = None, + ) -> ProcessResult: + active_context = context or ProcessingContext() + active_context.final = final + output: list[str] = [] + metadata: list[object] = [] + cursor = 0 + + while cursor < len(text): + probe = None + for rule in self._rules: + candidate = rule.probe(text, cursor, active_context) + if candidate.status != ProbeStatus.NO_MATCH: + probe = candidate + break + + if probe is None: + output.append(text[cursor]) + cursor += 1 + continue + + if probe.status == ProbeStatus.NEED_MORE: + break + if probe.consumed <= 0: + raise ValueError("A matching stream rule must consume source text") + + source = text[cursor : cursor + probe.consumed] + if probe.action == Action.PASS: + output.append( + probe.replacement if probe.replacement is not None else source + ) + elif probe.action in (Action.REPLACE, Action.UNWRAP): + output.append(probe.replacement or "") + elif probe.action == Action.DROP: + pass + else: + raise ValueError(f"Unsupported matching action: {probe.action}") + + for key in probe.state_deletes: + active_context.state.pop(key, None) + active_context.state.update(probe.state_updates) + metadata.extend(probe.metadata) + cursor += probe.consumed + + return ProcessResult( + text="".join(output), + metadata=tuple(metadata), + pending=text[cursor:], + consumed=cursor, + ) + + def feed(self, chunk: str) -> ProcessDelta: + self._source += chunk + self.context = ProcessingContext(state=deepcopy(self._initial_state)) + result = self.process(self._source, context=self.context) + if not result.text.startswith(self._emitted): + raise ValueError("A stream rule rewrote an already-emitted prefix") + delta = ProcessDelta( + text=result.text[len(self._emitted) :], + metadata=result.metadata[self._metadata_count :], + pending=result.pending, + ) + self._emitted = result.text + self._metadata_count = len(result.metadata) + return delta + + def finalize(self) -> ProcessDelta: + self.context = ProcessingContext(state=deepcopy(self._initial_state)) + result = self.process(self._source, final=True, context=self.context) + if not result.text.startswith(self._emitted): + raise ValueError("Finalization rewrote an already-emitted prefix") + delta = ProcessDelta( + text=result.text[len(self._emitted) :], + metadata=result.metadata[self._metadata_count :], + pending=result.pending, + ) + self._emitted = result.text + self._metadata_count = len(result.metadata) + return delta diff --git a/private_gpt/components/text_processing/models.py b/private_gpt/components/text_processing/models.py new file mode 100644 index 00000000..e98b3a5b --- /dev/null +++ b/private_gpt/components/text_processing/models.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum, auto +from typing import Any + + +class Action(Enum): + PASS = auto() + DROP = auto() + UNWRAP = auto() + REPLACE = auto() + HOLD = auto() + + +class ProbeStatus(Enum): + NO_MATCH = auto() + MATCH = auto() + NEED_MORE = auto() + + +@dataclass(frozen=True) +class ProbeResult: + status: ProbeStatus + consumed: int = 0 + action: Action = Action.PASS + replacement: str | None = None + metadata: tuple[Any, ...] = () + state_updates: dict[str, Any] = field(default_factory=dict) + state_deletes: tuple[str, ...] = () + + @classmethod + def no_match(cls) -> ProbeResult: + return cls(status=ProbeStatus.NO_MATCH) + + @classmethod + def need_more(cls) -> ProbeResult: + return cls(status=ProbeStatus.NEED_MORE, action=Action.HOLD) + + +@dataclass +class ProcessingContext: + final: bool = False + state: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ProcessResult: + text: str + metadata: tuple[Any, ...] + pending: str + consumed: int + + +@dataclass(frozen=True) +class ProcessDelta: + text: str + metadata: tuple[Any, ...] + pending: str diff --git a/private_gpt/components/text_processing/rules.py b/private_gpt/components/text_processing/rules.py new file mode 100644 index 00000000..ae9a56fe --- /dev/null +++ b/private_gpt/components/text_processing/rules.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import re +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Protocol + +from private_gpt.components.text_processing.models import ( + Action, + ProbeResult, + ProbeStatus, + ProcessingContext, +) + + +class StreamRule(Protocol): + name: str + priority: int + + def probe( + self, text: str, position: int, context: ProcessingContext + ) -> ProbeResult: ... + + +ResolveReferences = Callable[[list[str], ProcessingContext], list[Any]] +RenderReferences = Callable[[list[Any], ProcessingContext], tuple[str, tuple[Any, ...]]] + + +@dataclass +class DelimitedReferenceRule: + start_token: str + end_token: str + separator: str + resolve: ResolveReferences + render: RenderReferences + name: str = "delimited_reference" + priority: int = 100 + + def probe( + self, text: str, position: int, context: ProcessingContext + ) -> ProbeResult: + if not text.startswith(self.start_token, position): + return ProbeResult.no_match() + + end = text.find(self.end_token, position + len(self.start_token)) + if end == -1: + return ProbeResult.need_more() + + consumed = end + len(self.end_token) - position + identifiers = [ + identifier.strip() + for identifier in text[position + len(self.start_token) : end].split( + self.separator + ) + ] + references = self.resolve(identifiers, context) + if not references: + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=consumed, + action=Action.PASS, + ) + + replacement, metadata = self.render(references, context) + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=consumed, + action=Action.REPLACE, + replacement=replacement, + metadata=metadata, + ) + + +@dataclass +class BacktickUnwrapRule: + inner: StreamRule + name: str = "backtick_unwrap" + priority: int = 200 + code_state_key: str = "backtick_code_delimiter" + wrapper_state_key: str = "backtick_wrapper_delimiter" + + def probe( + self, text: str, position: int, context: ProcessingContext + ) -> ProbeResult: + if text[position] != "`": + return ProbeResult.no_match() + + delimiter_end = position + 1 + while delimiter_end < len(text) and text[delimiter_end] == "`": + delimiter_end += 1 + delimiter = text[position:delimiter_end] + + if context.state.get(self.wrapper_state_key) == delimiter: + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=len(delimiter), + action=Action.DROP, + state_deletes=(self.wrapper_state_key,), + ) + + if context.state.get(self.code_state_key) == delimiter: + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=len(delimiter), + action=Action.PASS, + state_deletes=(self.code_state_key,), + ) + + if delimiter_end == len(text): + if not context.final: + return ProbeResult.need_more() + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=len(delimiter), + action=Action.PASS, + ) + + inner_match = self.inner.probe(text, delimiter_end, context) + if inner_match.status == ProbeStatus.NEED_MORE: + return inner_match + if ( + inner_match.status == ProbeStatus.MATCH + and inner_match.action == Action.REPLACE + ): + consumed = len(delimiter) + inner_match.consumed + updates = dict(inner_match.state_updates) + deletes = list(inner_match.state_deletes) + if text.startswith(delimiter, position + consumed): + consumed += len(delimiter) + else: + updates[self.wrapper_state_key] = delimiter + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=consumed, + action=Action.UNWRAP, + replacement=inner_match.replacement, + metadata=inner_match.metadata, + state_updates=updates, + state_deletes=tuple(deletes), + ) + + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=len(delimiter), + action=Action.PASS, + state_updates={self.code_state_key: delimiter}, + ) + + +@dataclass +class LooseReferenceCleanupRule: + start_token: str + end_token: str + identifier_length: int + identifiers: tuple[str, ...] + name: str = "loose_reference_cleanup" + priority: int = 50 + + def __post_init__(self) -> None: + self._pattern = re.compile( + rf"{re.escape(self.start_token)}?[A-Z0-9]" + rf"{{{self.identifier_length}}}{re.escape(self.end_token)}?" + ) + + def probe( + self, text: str, position: int, context: ProcessingContext + ) -> ProbeResult: + match = self._pattern.match(text, position) + if match is None: + return ProbeResult.no_match() + word = match.group(0) + if word.startswith(self.start_token) and word.endswith(self.end_token): + return ProbeResult.no_match() + identifier = next( + (identifier for identifier in self.identifiers if identifier in word), + None, + ) + if identifier is None: + return ProbeResult.no_match() + replacement = word.replace(identifier, "", 1).strip() + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=len(word), + action=Action.REPLACE, + replacement=replacement, + ) diff --git a/tests/components/text_processing/test_engine.py b/tests/components/text_processing/test_engine.py new file mode 100644 index 00000000..e7a4c123 --- /dev/null +++ b/tests/components/text_processing/test_engine.py @@ -0,0 +1,131 @@ +from dataclasses import dataclass + +from private_gpt.components.text_processing import ( + Action, + IncrementalTextProcessor, + ProbeResult, + ProbeStatus, + ProcessingContext, +) + + +@dataclass +class TokenRule: + token: str + action: Action + replacement: str | None = None + priority: int = 100 + name: str = "token" + + def probe( + self, text: str, position: int, context: ProcessingContext + ) -> ProbeResult: + remaining = text[position:] + if self.token.startswith(remaining) and len(remaining) < len(self.token): + return ProbeResult.need_more() + if not text.startswith(self.token, position): + return ProbeResult.no_match() + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=len(self.token), + action=self.action, + replacement=self.replacement, + ) + + +def test_processor_passes_literal_text_without_rules() -> None: + result = IncrementalTextProcessor([]).process("plain text") + + assert result.text == "plain text" + assert result.pending == "" + + +def test_processor_supports_pass_drop_unwrap_and_replace() -> None: + processor = IncrementalTextProcessor( + [ + TokenRule("", Action.PASS), + TokenRule("", Action.DROP), + TokenRule("", Action.UNWRAP, "body"), + TokenRule("", Action.REPLACE, "replacement"), + ] + ) + + result = processor.process("|||") + + assert result.text == "||body|replacement" + + +def test_need_more_holds_only_the_unresolved_suffix() -> None: + processor = IncrementalTextProcessor( + [TokenRule("", Action.REPLACE, "done")] + ) + + result = processor.process("safe None: + processor = IncrementalTextProcessor( + [TokenRule("", Action.REPLACE, "done")] + ) + + result = processor.process("safe None: + processor = IncrementalTextProcessor( + [ + TokenRule("token", Action.REPLACE, "low", priority=10), + TokenRule("token", Action.REPLACE, "high", priority=20), + ] + ) + + assert processor.process("token").text == "high" + + +def test_feed_emits_only_new_safe_text() -> None: + processor = IncrementalTextProcessor( + [TokenRule("", Action.REPLACE, "done")] + ) + + first = processor.feed("before after") + + assert first.text == "before " + assert first.pending == " ProbeResult: + if not text.startswith("#", position): + return ProbeResult.no_match() + count = int(context.state.get("count", 0)) + 1 + return ProbeResult( + status=ProbeStatus.MATCH, + consumed=1, + action=Action.REPLACE, + replacement=str(count), + state_updates={"count": count}, + ) + + +def test_feed_reparses_from_clean_initial_rule_state() -> None: + processor = IncrementalTextProcessor([StatefulRule()]) + + assert processor.feed("#").text == "1" + assert processor.feed("#").text == "2" + assert processor.context.state == {"count": 2} diff --git a/tests/engines/test_citation_chunk_boundaries.py b/tests/engines/test_citation_chunk_boundaries.py index 7d4df9fd..ad677fa4 100644 --- a/tests/engines/test_citation_chunk_boundaries.py +++ b/tests/engines/test_citation_chunk_boundaries.py @@ -7,7 +7,6 @@ from private_gpt.components.engines.citations.utils import ( extract_citations_by_original_text, ) - TEXT = ( "Format: `[XXXX]`. Correct: `[AB12]`. " "Invalid: `[[AB12]]`, `(AB12)`. " diff --git a/tests/engines/test_citation_parser_robustness.py b/tests/engines/test_citation_parser_robustness.py index 76ccfbd8..a7d1bddb 100644 --- a/tests/engines/test_citation_parser_robustness.py +++ b/tests/engines/test_citation_parser_robustness.py @@ -118,10 +118,6 @@ def test_incomplete_citation_is_withheld_with_no_false_citation() -> None: 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]." @@ -132,8 +128,7 @@ def test_placeholder_like_model_output_does_not_capture_real_citation() -> None: ) assert formatted == ( - "Literal \ue000citationn0\ue001 then " - f"{format_cite(0, document, 0)}." + f"Literal \ue000citationn0\ue001 then {format_cite(0, document, 0)}." ) assert len(citations) == 1 diff --git a/tests/engines/test_processor_citations.py b/tests/engines/test_processor_citations.py index baa6db07..9b76a8b1 100644 --- a/tests/engines/test_processor_citations.py +++ b/tests/engines/test_processor_citations.py @@ -181,10 +181,10 @@ async def test_process_citations_trailing_backtick_with_stop_event() -> None: # start, delta (without backtick), delta (with backtick), stop assert len(result) == 4 - + combined = "" for r in result: if isinstance(r, RawContentBlockDeltaEvent) and isinstance(r.delta, TextDelta): combined += r.delta.text or "" - + assert combined == "This is a test with a trailing backtick `"