github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-16 17:00:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: tick quotes

Browse Source
This commit is contained in:
Javier Martinez
2026-07-16 15:48:17 +02:00
parent 4aec14152f
commit 2b75062d38
5 changed files with 599 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

177
private_gpt/components/engines/citations/utils.py
View File

@@ -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

118
tests/engines/test_citation_chunk_boundaries.py Normal file
View File

@@ -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 == []

155
tests/engines/test_citation_parser_robustness.py Normal file
View File

@@ -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}

43
tests/engines/test_citation_streaming_quotes.py Normal file
View File

@@ -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: <citation id='LLFB'" in emitted_text
assert ", <citation id='CPCY'" in emitted_text
assert "</citation>`. End." not in emitted_text

185
tests/engines/test_processor_citations_streaming_robustness.py Normal file
View File

@@ -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