From 833982ac814a70156d98b5e901b41edfc8365266 Mon Sep 17 00:00:00 2001 From: TommasoMoroHtx Date: Thu, 5 Jun 2025 11:08:49 +0200 Subject: [PATCH] fix(docstore): preserve retrieval ranking order in lancedb get() (#745) --- libs/kotaemon/kotaemon/storages/docstores/lancedb.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/libs/kotaemon/kotaemon/storages/docstores/lancedb.py b/libs/kotaemon/kotaemon/storages/docstores/lancedb.py index fa34a747..36ddb69f 100644 --- a/libs/kotaemon/kotaemon/storages/docstores/lancedb.py +++ b/libs/kotaemon/kotaemon/storages/docstores/lancedb.py @@ -113,14 +113,17 @@ class LanceDBDocumentStore(BaseDocumentStore): ) except (ValueError, FileNotFoundError): docs = [] - return [ - Document( - id_=doc["id"], + + # return the documents using the order of original ids (which were ordered by score) + doc_dict = { + doc["id"]: Document( + d_=doc["id"], text=doc["text"] if doc["text"] else "", metadata=json.loads(doc["attributes"]), ) for doc in docs - ] + } + return [doc_dict[_id] for _id in ids if _id in doc_dict] def delete(self, ids: Union[List[str], str], refresh_indices: bool = True): """Delete document by id"""