mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-12 04:40:09 +02:00
24 lines
648 B
Python
24 lines
648 B
Python
from collections.abc import Sequence
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from onyx.chat.models import LlmDoc
|
|
from onyx.context.search.models import InferenceChunk
|
|
|
|
|
|
class DocumentIdOrderMapping(BaseModel):
|
|
order_mapping: dict[str, int]
|
|
|
|
|
|
def map_document_id_order(
|
|
chunks: Sequence[InferenceChunk | LlmDoc], one_indexed: bool = True
|
|
) -> DocumentIdOrderMapping:
|
|
order_mapping = {}
|
|
current = 1 if one_indexed else 0
|
|
for chunk in chunks:
|
|
if chunk.document_id not in order_mapping:
|
|
order_mapping[chunk.document_id] = current
|
|
current += 1
|
|
|
|
return DocumentIdOrderMapping(order_mapping=order_mapping)
|