roadmap(1.3): Update citation prompt to use new whole document structure

- Modified the document extraction and citation formatting to accommodate a new structure that includes a `chunks` list for each document.
- Enhanced the citation format to reference `chunk_id` instead of `source_id`, ensuring accurate citations in the UI.
- Updated various components, including the connector service and reranker service, to handle the new document format and maintain compatibility with existing functionalities.
- Improved documentation and comments to reflect changes in the data structure and citation requirements.
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2025-12-14 22:07:31 -08:00
parent ed6fc10133
commit fea1837186
9 changed files with 1054 additions and 1122 deletions

File diff suppressed because it is too large Load diff

View file

@ -22,14 +22,18 @@ class RerankerService:
self, query_text: str, documents: list[dict[str, Any]]
) -> list[dict[str, Any]]:
"""
Rerank documents using the configured reranker
Rerank documents using the configured reranker.
Documents can be either:
- Document-grouped (new format): Has `document_id`, `chunks` list, and `content` (concatenated)
- Chunk-based (legacy format): Individual chunks with `chunk_id` and `content`
Args:
query_text: The query text to use for reranking
documents: List of document dictionaries to rerank
Returns:
List[Dict[str, Any]]: Reranked documents
List[Dict[str, Any]]: Reranked documents with preserved structure
"""
if not self.reranker_instance or not documents:
return documents
@ -38,7 +42,9 @@ class RerankerService:
# Create Document objects for the rerankers library
reranker_docs = []
for i, doc in enumerate(documents):
chunk_id = doc.get("chunk_id", f"chunk_{i}")
# Use document_id for matching
doc_id = doc.get("document_id") or f"doc_{i}"
# Use concatenated content for reranking
content = doc.get("content", "")
score = doc.get("score", 0.0)
document_info = doc.get("document", {})
@ -46,12 +52,14 @@ class RerankerService:
reranker_docs.append(
RerankerDocument(
text=content,
doc_id=chunk_id,
doc_id=doc_id,
metadata={
"document_id": document_info.get("id", ""),
"document_title": document_info.get("title", ""),
"document_type": document_info.get("document_type", ""),
"rrf_score": score,
# Track original index for fallback matching
"original_index": i,
},
)
)
@ -62,21 +70,33 @@ class RerankerService:
)
# Process the results from the reranker
# Convert to serializable dictionaries
# Convert to serializable dictionaries while preserving full structure
serialized_results = []
for result in reranking_results.results:
# Find the original document by id
original_doc = next(
(
doc
for doc in documents
if doc.get("chunk_id") == result.document.doc_id
),
None,
)
result_doc_id = result.document.doc_id
original_index = result.document.metadata.get("original_index")
# Find the original document by document_id
original_doc = None
for doc in documents:
if doc.get("document_id") == result_doc_id:
original_doc = doc
break
# Fallback to original index if ID matching fails
if (
original_doc is None
and original_index is not None
and 0 <= original_index < len(documents)
):
original_doc = documents[original_index]
if original_doc:
# Create a new document with the reranked score
# Create a deep copy to preserve the full structure including chunks
reranked_doc = original_doc.copy()
# Preserve chunks list if present (important for citation formatting)
if "chunks" in original_doc:
reranked_doc["chunks"] = original_doc["chunks"]
reranked_doc["score"] = float(result.score)
reranked_doc["rank"] = result.rank
serialized_results.append(reranked_doc)