cognee-community-vector-adapter-qdrant 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

cognee_community_vector_adapter_qdrant/QdrantDatasetDatabaseHandler.py

@@ -0,0 +1,39 @@
+from typing import Optional
+from uuid import UUID
+
+from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
+from cognee.infrastructure.databases.vector import get_vectordb_config
+from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
+from cognee.modules.users.models import DatasetDatabase, User
+
+
+class QdrantDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    @classmethod
+    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
+        vector_config = get_vectordb_config()
+
+        if vector_config.vector_db_provider != "qdrant":
+            raise ValueError(
+                "QdrantDatasetDatabaseHandler can only be used with the"
+                "Qdrant vector database provider."
+            )
+
+        vector_db_name = f"{dataset_id}"
+
+        return {
+            "vector_database_provider": vector_config.vector_db_provider,
+            "vector_database_url": vector_config.vector_db_url,
+            "vector_database_key": vector_config.vector_db_key,
+            "vector_database_name": vector_db_name,
+            "vector_dataset_database_handler": "qdrant",
+        }
+
+    @classmethod
+    async def delete_dataset(cls, dataset_database: DatasetDatabase) -> None:
+        vector_engine = create_vector_engine(
+            vector_db_provider=dataset_database.vector_database_provider,
+            vector_db_url=dataset_database.vector_database_url,
+            vector_db_key=dataset_database.vector_database_key,
+            vector_db_name=dataset_database.vector_database_name,
+        )
+        await vector_engine.prune()

cognee_community_vector_adapter_qdrant/qdrant_adapter.py

@@ -45,8 +45,16 @@ class QDrantAdapter(VectorDBInterface):
     api_key: str = None
     qdrant_path: str = None
 
-    def __init__(self, url, api_key, embedding_engine: EmbeddingEngine, qdrant_path=None):
+    def __init__(
+        self,
+        url,
+        api_key,
+        embedding_engine: EmbeddingEngine,
+        qdrant_path=None,
+        database_name: str = "cognee_db",
+    ):
         self.embedding_engine = embedding_engine
+        self.database_name = database_name
 
         if qdrant_path is not None:
             self.qdrant_path = qdrant_path
@@ -86,9 +94,24 @@ class QDrantAdapter(VectorDBInterface):
                     vectors_config={
                         "text": models.VectorParams(
                             size=self.embedding_engine.get_vector_size(),
-                            distance="Cosine",
+                            distance=models.Distance.COSINE,
                         )
                     },
+                    # With this config definition, we avoid creating a global index
+                    hnsw_config=models.HnswConfigDiff(
+                        payload_m=16,
+                        m=0,
+                    ),
+                )
+                # This index co-locates vectors from the same dataset together,
+                # which can improve performance
+                await client.create_payload_index(
+                    collection_name=collection_name,
+                    field_name="database_name",
+                    field_schema=models.KeywordIndexParams(
+                        type=models.KeywordIndexType.KEYWORD,
+                        is_tenant=True,
+                    ),
                 )
 
             await client.close()
@@ -105,7 +128,7 @@ class QDrantAdapter(VectorDBInterface):
         def convert_to_qdrant_point(data_point: DataPoint):
             return models.PointStruct(
                 id=str(data_point.id),
-                payload=data_point.model_dump(),
+                payload={**data_point.model_dump(), "database_name": self.database_name},
                 vector={"text": data_vectors[data_points.index(data_point)]},
             )
 
@@ -157,6 +180,7 @@ class QDrantAdapter(VectorDBInterface):
         query_vector: list[float] | None = None,
         limit: int | None = 15,
         with_vector: bool = False,
+        include_payload: bool = False,
     ) -> list[ScoredResult]:
         if query_text is None and query_vector is None:
             raise MissingQueryParameterError()
@@ -182,9 +206,20 @@ class QDrantAdapter(VectorDBInterface):
             query_result = await client.query_points(
                 collection_name=collection_name,
                 query=query_vector,
+                query_filter=models.Filter(
+                    must=[
+                        models.FieldCondition(
+                            key="database_name",
+                            match=models.MatchValue(
+                                value=self.database_name,
+                            ),
+                        )
+                    ]
+                ),
                 using="text",
                 limit=limit,
                 with_vectors=with_vector,
+                with_payload=include_payload,
             )
 
             await client.close()
@@ -195,11 +230,13 @@ class QDrantAdapter(VectorDBInterface):
             return [
                 ScoredResult(
                     id=parse_id(str(result.id)),
-                    payload={
+                    payload=None
+                    if not result.payload
+                    else {
                         **result.payload,
                         "id": parse_id(str(result.id)),
                     },
-                    score=1 - result.score if hasattr(result, 'score') else 1.0,
+                    score=1 - result.score if hasattr(result, "score") else 1.0,
                 )
                 for result in results
             ]
@@ -215,6 +252,7 @@ class QDrantAdapter(VectorDBInterface):
         query_texts: list[str],
         limit: int | None = None,
         with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         """
         Perform batch search in a Qdrant collection with dynamic search requests.
@@ -225,6 +263,7 @@ class QDrantAdapter(VectorDBInterface):
         - limit (int): List of result limits for search requests.
         - with_vectors (bool, optional): Bool indicating whether to return
                                          vectors for search requests.
+        - include_payload (bool, optional): Bool indicating whether to return payload in results.
 
         Returns:
         - results: The search results from Qdrant.
@@ -246,8 +285,19 @@ class QDrantAdapter(VectorDBInterface):
             query_results = await client.query_batch(
                 collection_name=collection_name,
                 query_texts=query_texts,
+                query_filter=models.Filter(
+                    must=[
+                        models.FieldCondition(
+                            key="database_name",
+                            match=models.MatchValue(
+                                value=self.database_name,
+                            ),
+                        )
+                    ]
+                ),
                 limit=limit,
                 with_vectors=with_vectors,
+                with_payload=include_payload,
             )
 
             await client.close()
@@ -255,10 +305,9 @@ class QDrantAdapter(VectorDBInterface):
             # Extract points from each query result and filter by score
             filtered_results = []
             for query_result in query_results:
-                points = query_result.points if hasattr(query_result, 'points') else []
+                points = query_result.points if hasattr(query_result, "points") else []
                 filtered_points = [
-                    result for result in points
-                    if hasattr(result, 'score') and result.score > 0.9
+                    result for result in points if hasattr(result, "score") and result.score > 0.9
                 ]
                 filtered_results.append(filtered_points)
 
@@ -279,7 +328,22 @@ class QDrantAdapter(VectorDBInterface):
         response = await client.get_collections()
 
         for collection in response.collections:
-            await client.delete_collection(collection.name)
+            await client.delete(
+                collection.name,
+                points_selector=models.FilterSelector(
+                    filter=models.Filter(
+                        must=[
+                            models.FieldCondition(
+                                key="database_name",
+                                match=models.MatchValue(value=self.database_name),
+                            )
+                        ]
+                    )
+                ),
+            )
+            remaining_points = await client.count(collection_name=collection.name)
+            if remaining_points.count == 0:
+                await client.delete_collection(collection_name=collection.name)
 
         await client.close()
 
@@ -295,7 +359,23 @@ class QDrantAdapter(VectorDBInterface):
 
         response = await client.get_collections()
 
-        result = [collection.name for collection in response.collections]
+        # We do this filtering because one user could see another user's collections otherwise
+        result = []
+        for collection in response.collections:
+            relevant_count = await client.count(
+                collection_name=collection.name,
+                count_filter=models.Filter(
+                    must=[
+                        models.FieldCondition(
+                            key="database_name", match=models.MatchValue(value=self.database_name)
+                        )
+                    ]
+                ),
+                exact=True,
+            )
+
+            if relevant_count.count > 0:
+                result.append(collection.name)
 
         await client.close()
 

cognee_community_vector_adapter_qdrant/register.py

@@ -1,5 +1,8 @@
+from cognee.infrastructure.databases.dataset_database_handler import use_dataset_database_handler
 from cognee.infrastructure.databases.vector import use_vector_adapter
 
 from .qdrant_adapter import QDrantAdapter
+from .QdrantDatasetDatabaseHandler import QdrantDatasetDatabaseHandler
 
 use_vector_adapter("qdrant", QDrantAdapter)
+use_dataset_database_handler("qdrant", QdrantDatasetDatabaseHandler, "qdrant")

{cognee_community_vector_adapter_qdrant-0.1.0.dist-info → cognee_community_vector_adapter_qdrant-0.2.1.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: cognee-community-vector-adapter-qdrant
-Version: 0.1.0
+Version: 0.2.1
 Summary: Qdrant vector database adapter for cognee
 Requires-Python: >=3.11,<=3.13
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Requires-Dist: cognee (>=0.4.0)
+Requires-Dist: cognee (==0.5.2)
 Requires-Dist: instructor (>=1.11)
 Requires-Dist: qdrant-client (>=1.16.0)
 Requires-Dist: starlette (>=0.48.0)
@@ -49,6 +49,11 @@ Import and register the adapter in your code:
 from cognee_community_vector_adapter_qdrant import register
 ```
 
+Also, specify the dataset handler in the .env file:
+```dotenv
+VECTOR_DATASET_DATABASE_HANDLER="qdrant"
+```
+
 ## Example
 See example in `example.py` file.
 

cognee_community_vector_adapter_qdrant-0.2.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+cognee_community_vector_adapter_qdrant/QdrantDatasetDatabaseHandler.py,sha256=mQFgATAdZeiBayoo4poLsXnCVrFPTLN3zSJdGzq0brk,1668
+cognee_community_vector_adapter_qdrant/__init__.py,sha256=OYsRMTyBnNuusKB3rkQmwCppfdEf__AuXMw_Ma1k6lQ,71
+cognee_community_vector_adapter_qdrant/qdrant_adapter.py,sha256=ONpLcjtufnSQEJoXLxfk77oOM2idiU3eGkZWSa5avFo,13591
+cognee_community_vector_adapter_qdrant/register.py,sha256=_CUgC1Ton9HfPfsyi6dEQ5H1J47bXHAwFHIceWSo4SI,406
+cognee_community_vector_adapter_qdrant-0.2.1.dist-info/METADATA,sha256=QY87_e1VZ73U0ACxk6ba4YsCJXmsKZRP-YJUD3bbWYI,1645
+cognee_community_vector_adapter_qdrant-0.2.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+cognee_community_vector_adapter_qdrant-0.2.1.dist-info/RECORD,,

cognee_community_vector_adapter_qdrant-0.1.0.dist-info/RECORD

@@ -1,6 +0,0 @@
-cognee_community_vector_adapter_qdrant/__init__.py,sha256=OYsRMTyBnNuusKB3rkQmwCppfdEf__AuXMw_Ma1k6lQ,71
-cognee_community_vector_adapter_qdrant/qdrant_adapter.py,sha256=hFMmSaUArnhTZHR8I9VW3fHdmO-1VHw8TD9-fJ4rs-c,10379
-cognee_community_vector_adapter_qdrant/register.py,sha256=K0cIQGN3an79wWCMXAIgwsymkloHGV2_joy7G-4aiB8,158
-cognee_community_vector_adapter_qdrant-0.1.0.dist-info/METADATA,sha256=0dcxsPnOFHs-gDGnrmFJR9up07mdE6AAQ7tkvRhnPbA,1537
-cognee_community_vector_adapter_qdrant-0.1.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-cognee_community_vector_adapter_qdrant-0.1.0.dist-info/RECORD,,