72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""Utilities for accumulating RAG source citations across multiple searches."""
|
|
|
|
import json
|
|
from typing import Any, Dict
|
|
|
|
|
|
def _citation_item_key(item: Any) -> str:
|
|
"""Build a stable identity for one RAG citation chunk."""
|
|
if not isinstance(item, dict):
|
|
return json.dumps(item, ensure_ascii=False, sort_keys=True, default=str)
|
|
|
|
chunk_id = item.get("id")
|
|
if chunk_id not in (None, ""):
|
|
return f"id:{chunk_id}"
|
|
|
|
location = {
|
|
"file_id": item.get("file_id"),
|
|
"page_idx": item.get("page_idx"),
|
|
"positions": item.get("positions"),
|
|
}
|
|
if any(value not in (None, "", []) for value in location.values()):
|
|
return "location:" + json.dumps(
|
|
location, ensure_ascii=False, sort_keys=True, default=str
|
|
)
|
|
|
|
# index/score may differ between rewritten queries and do not identify a source.
|
|
fallback = {
|
|
key: value
|
|
for key, value in item.items()
|
|
if key not in {"index", "score", "text"}
|
|
}
|
|
return "fallback:" + json.dumps(
|
|
fallback, ensure_ascii=False, sort_keys=True, default=str
|
|
)
|
|
|
|
|
|
def merge_rag_source_citations(
|
|
accumulated: Dict[str, Any], incoming: Any
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Append a RAG call's citations without losing earlier calls or duplicating chunks.
|
|
|
|
Citation text is intentionally removed here because the stream response has never
|
|
exposed it; the full text remains available inside the workflow for answer generation.
|
|
"""
|
|
if not isinstance(incoming, dict) or not incoming:
|
|
return accumulated
|
|
|
|
for doc_name, chunks in incoming.items():
|
|
if isinstance(chunks, list):
|
|
target = accumulated.setdefault(doc_name, [])
|
|
if not isinstance(target, list):
|
|
# Be defensive about malformed/legacy data while preserving both values.
|
|
target = [target]
|
|
accumulated[doc_name] = target
|
|
|
|
seen = {_citation_item_key(item) for item in target}
|
|
for item in chunks:
|
|
sanitized_item = (
|
|
{key: value for key, value in item.items() if key != "text"}
|
|
if isinstance(item, dict)
|
|
else item
|
|
)
|
|
item_key = _citation_item_key(sanitized_item)
|
|
if item_key not in seen:
|
|
target.append(sanitized_item)
|
|
seen.add(item_key)
|
|
elif doc_name not in accumulated:
|
|
accumulated[doc_name] = chunks
|
|
|
|
return accumulated
|