langroid 0.1.139__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/vector_store/qdrantdb.py

@@ -1,6 +1,9 @@
+import hashlib
+import json
 import logging
 import os
-from typing import List, Optional, Sequence, Tuple
+import uuid
+from typing import List, Optional, Sequence, Tuple, TypeVar
 
 from dotenv import load_dotenv
 from qdrant_client import QdrantClient
@@ -26,6 +29,33 @@ from langroid.vector_store.base import VectorStore, VectorStoreConfig
 logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T")
+
+
+def from_optional(x: Optional[T], default: T) -> T:
+    if x is None:
+        return default
+
+    return x
+
+
+def is_valid_uuid(uuid_to_test: str) -> bool:
+    """
+    Check if a given string is a valid UUID.
+    """
+    try:
+        uuid_obj = uuid.UUID(uuid_to_test)
+        return str(uuid_obj) == uuid_to_test
+    except Exception:
+        pass
+    # Check for valid unsigned 64-bit integer
+    try:
+        int_value = int(uuid_to_test)
+        return 0 <= int_value <= 18446744073709551615
+    except ValueError:
+        return False
+
+
 class QdrantDBConfig(VectorStoreConfig):
     cloud: bool = True
     collection_name: str | None = "temp"
@@ -35,7 +65,7 @@ class QdrantDBConfig(VectorStoreConfig):
 
 
 class QdrantDB(VectorStore):
-    def __init__(self, config: QdrantDBConfig):
+    def __init__(self, config: QdrantDBConfig = QdrantDBConfig()):
         super().__init__(config)
         self.config = config
         emb_model = EmbeddingModel.create(config.embedding)
@@ -112,8 +142,10 @@ class QdrantDB(VectorStore):
         n_non_empty_deletes = 0
         for name in coll_names:
             info = self.client.get_collection(collection_name=name)
-            n_empty_deletes += (info.points_count or 0) == 0
-            n_non_empty_deletes += (info.points_count or 0) > 0
+            points_count = from_optional(info.points_count, 0)
+
+            n_empty_deletes += points_count == 0
+            n_non_empty_deletes += points_count > 0
             self.client.delete_collection(collection_name=name)
         logger.warning(
             f"""
@@ -138,7 +170,12 @@ class QdrantDB(VectorStore):
         for coll in colls:
             try:
                 counts.append(
-                    self.client.get_collection(collection_name=coll.name).points_count
+                    from_optional(
+                        self.client.get_collection(
+                            collection_name=coll.name
+                        ).points_count,
+                        0,
+                    )
                 )
             except Exception:
                 logger.warning(f"Error getting collection {coll.name}")
@@ -158,7 +195,10 @@ class QdrantDB(VectorStore):
         collections = self.list_collections()
         if collection_name in collections:
             coll = self.client.get_collection(collection_name=collection_name)
-            if coll.status == CollectionStatus.GREEN and (coll.points_count or 0) > 0:
+            if (
+                coll.status == CollectionStatus.GREEN
+                and from_optional(coll.points_count, 0) > 0
+            ):
                 logger.warning(f"Non-empty Collection {collection_name} already exists")
                 if not replace:
                     logger.warning("Not replacing collection")
@@ -182,10 +222,15 @@ class QdrantDB(VectorStore):
             logger.setLevel(level)
 
     def add_documents(self, documents: Sequence[Document]) -> None:
+        # Add id to metadata if not already present
         super().maybe_add_ids(documents)
+        # Fix the ids due to qdrant finickiness
+        for doc in documents:
+            doc.metadata.id = str(self._to_int_or_uuid(doc.metadata.id))
         colls = self.list_collections(empty=True)
         if len(documents) == 0:
             return
+        document_dicts = [doc.dict() for doc in documents]
         embedding_vecs = self.embedding_fn([doc.content for doc in documents])
         if self.config.collection_name is None:
             raise ValueError("No collection name set, cannot ingest docs")
@@ -201,7 +246,7 @@ class QdrantDB(VectorStore):
                 points=Batch(
                     ids=ids[i : i + b],
                     vectors=embedding_vecs[i : i + b],
-                    payloads=documents[i : i + b],
+                    payloads=document_dicts[i : i + b],
                 ),
             )
 
@@ -210,19 +255,42 @@ class QdrantDB(VectorStore):
 
     def _to_int_or_uuid(self, id: str) -> int | str:
         try:
-            return int(id)
+            int_val = int(id)
+            if is_valid_uuid(id):
+                return int_val
         except ValueError:
+            pass
+
+        # If doc_id is already a valid UUID, return it as is
+        if isinstance(id, str) and is_valid_uuid(id):
             return id
 
-    def get_all_documents(self) -> List[Document]:
+        # Otherwise, generate a UUID from the doc_id
+        # Convert doc_id to string if it's not already
+        id_str = str(id)
+
+        # Hash the document ID using SHA-1
+        hash_object = hashlib.sha1(id_str.encode())
+        hash_digest = hash_object.hexdigest()
+
+        # Truncate or manipulate the hash to fit into a UUID (128 bits)
+        uuid_str = hash_digest[:32]
+
+        # Format this string into a UUID format
+        formatted_uuid = uuid.UUID(uuid_str)
+
+        return str(formatted_uuid)
+
+    def get_all_documents(self, where: str = "") -> List[Document]:
         if self.config.collection_name is None:
             raise ValueError("No collection name set, cannot retrieve docs")
         docs = []
         offset = 0
+        filter = Filter() if where == "" else Filter.parse_obj(json.loads(where))
         while True:
             results, next_page_offset = self.client.scroll(
                 collection_name=self.config.collection_name,
-                scroll_filter=None,
+                scroll_filter=filter,
                 offset=offset,
                 limit=10_000,  # try getting all at once, if not we keep paging
                 with_payload=True,
@@ -260,7 +328,10 @@ class QdrantDB(VectorStore):
     ) -> List[Tuple[Document, float]]:
         embedding = self.embedding_fn([text])[0]
         # TODO filter may not work yet
-        filter = Filter() if where is None else Filter.from_json(where)  # type: ignore
+        if where is None or where == "":
+            filter = Filter()
+        else:
+            filter = Filter.parse_obj(json.loads(where))
         if self.config.collection_name is None:
             raise ValueError("No collection name set, cannot search")
         search_result: List[ScoredPoint] = self.client.search(