From 8f7031b882d5c3581fd6bae14cfef83cfbed5191 Mon Sep 17 00:00:00 2001 From: yangdx Date: Fri, 15 Aug 2025 16:46:52 +0800 Subject: [PATCH] Add get_vectors_by_ids method to QdrantVectorDBStorage --- lightrag/kg/qdrant_impl.py | 44 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/lightrag/kg/qdrant_impl.py b/lightrag/kg/qdrant_impl.py index 1686cba6..4ece163c 100644 --- a/lightrag/kg/qdrant_impl.py +++ b/lightrag/kg/qdrant_impl.py @@ -402,6 +402,50 @@ class QdrantVectorDBStorage(BaseVectorStorage): ) return [] + async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]: + """Get vectors by their IDs, returning only ID and vector data for efficiency + + Args: + ids: List of unique identifiers + + Returns: + Dictionary mapping IDs to their vector embeddings + Format: {id: [vector_values], ...} + """ + if not ids: + return {} + + try: + # Convert to Qdrant compatible IDs + qdrant_ids = [compute_mdhash_id_for_qdrant(id) for id in ids] + + # Retrieve the points by IDs with vectors + results = self._client.retrieve( + collection_name=self.final_namespace, + ids=qdrant_ids, + with_vectors=True, # Important: request vectors + with_payload=True, + ) + + vectors_dict = {} + for point in results: + if point and point.vector is not None and point.payload: + # Get original ID from payload + original_id = point.payload.get("id") + if original_id: + # Convert numpy array to list if needed + vector_data = point.vector + if isinstance(vector_data, np.ndarray): + vector_data = vector_data.tolist() + vectors_dict[original_id] = vector_data + + return vectors_dict + except Exception as e: + logger.error( + f"[{self.workspace}] Error retrieving vectors by IDs from {self.namespace}: {e}" + ) + return {} + async def drop(self) -> dict[str, str]: """Drop all vector data from storage and clean up resources