mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-06-15 18:40:56 +02:00
29f5f4edfa
* removed some dead code and fixed citations when a search request is made with sections selected * fix black formatting issue
25 lines
703 B
Python
25 lines
703 B
Python
from collections.abc import Sequence
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from onyx.chat.models import LlmDoc
|
|
from onyx.chat.models import OnyxContext
|
|
from onyx.context.search.models import InferenceChunk
|
|
|
|
|
|
class DocumentIdOrderMapping(BaseModel):
|
|
order_mapping: dict[str, int]
|
|
|
|
|
|
def map_document_id_order(
|
|
chunks: Sequence[InferenceChunk | LlmDoc | OnyxContext], 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)
|