cognee-community-vector-adapter-qdrant 0.0.3__tar.gz → 0.2.0__tar.gz

{cognee_community_vector_adapter_qdrant-0.0.3 → cognee_community_vector_adapter_qdrant-0.2.0}/PKG-INFO

@@ -1,14 +1,16 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: cognee-community-vector-adapter-qdrant
-Version: 0.0.3
+Version: 0.2.0
 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.2.4)
-Requires-Dist: qdrant-client (>=1.14.2)
+Requires-Dist: cognee (==0.5.1)
+Requires-Dist: instructor (>=1.11)
+Requires-Dist: qdrant-client (>=1.16.0)
+Requires-Dist: starlette (>=0.48.0)
 Description-Content-Type: text/markdown
 
 # Cognee Qdrant Adapter
@@ -47,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.0.3 → cognee_community_vector_adapter_qdrant-0.2.0}/README.md

@@ -34,5 +34,10 @@ 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.0/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-0.0.3 → cognee_community_vector_adapter_qdrant-0.2.0}/cognee_community_vector_adapter_qdrant/qdrant_adapter.py

@@ -1,18 +1,16 @@
 import asyncio
-from typing import Dict, List, Optional
-from qdrant_client import AsyncQdrantClient, models
-
-from cognee.shared.logging_utils import get_logger
 
-from cognee.infrastructure.engine import DataPoint
-from cognee.infrastructure.engine.utils import parse_id
 from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
 from cognee.infrastructure.databases.vector import VectorDBInterface
-from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import (
     EmbeddingEngine,
 )
 from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
+from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
+from cognee.infrastructure.engine import DataPoint
+from cognee.infrastructure.engine.utils import parse_id
+from cognee.shared.logging_utils import get_logger
+from qdrant_client import AsyncQdrantClient, models
 
 logger = get_logger("QDrantAdapter")
 
@@ -23,19 +21,19 @@ class IndexSchema(DataPoint):
     metadata: dict = {"index_fields": ["text"]}
 
 
-def create_hnsw_config(hnsw_config: Dict):
+def create_hnsw_config(hnsw_config: dict):
     if hnsw_config is not None:
         return models.HnswConfig()
     return None
 
 
-def create_optimizers_config(optimizers_config: Dict):
+def create_optimizers_config(optimizers_config: dict):
     if optimizers_config is not None:
         return models.OptimizersConfig()
     return None
 
 
-def create_quantization_config(quantization_config: Dict):
+def create_quantization_config(quantization_config: dict):
     if quantization_config is not None:
         return models.QuantizationConfig()
     return None
@@ -48,9 +46,15 @@ class QDrantAdapter(VectorDBInterface):
     qdrant_path: str = None
 
     def __init__(
-        self, url, api_key, embedding_engine: EmbeddingEngine, qdrant_path=None
+        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
@@ -67,7 +71,7 @@ class QDrantAdapter(VectorDBInterface):
 
         return AsyncQdrantClient(location=":memory:")
 
-    async def embed_data(self, data: List[str]) -> List[float]:
+    async def embed_data(self, data: list[str]) -> list[float]:
         return await self.embedding_engine.embed_text(data)
 
     async def has_collection(self, collection_name: str) -> bool:
@@ -90,16 +94,29 @@ 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()
 
-    async def create_data_points(
-        self, collection_name: str, data_points: List[DataPoint]
-    ):
+    async def create_data_points(self, collection_name: str, data_points: list[DataPoint]):
         from qdrant_client.http.exceptions import UnexpectedResponse
 
         client = self.get_qdrant_client()
@@ -111,14 +128,15 @@ 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)]},
             )
 
         points = [convert_to_qdrant_point(point) for point in data_points]
 
         try:
-            client.upload_points(collection_name=collection_name, points=points)
+            # Use upsert for AsyncQdrantClient (upload_points doesn't exist or is sync)
+            await client.upsert(collection_name=collection_name, points=points)
         except UnexpectedResponse as error:
             if "Collection not found" in str(error):
                 raise CollectionNotFoundError(
@@ -151,22 +169,18 @@ class QDrantAdapter(VectorDBInterface):
 
     async def retrieve(self, collection_name: str, data_point_ids: list[str]):
         client = self.get_qdrant_client()
-        results = await client.retrieve(
-            collection_name, data_point_ids, with_payload=True
-        )
+        results = await client.retrieve(collection_name, data_point_ids, with_payload=True)
         await client.close()
         return results
 
     async def search(
         self,
         collection_name: str,
-        query_text: Optional[str] = None,
-        query_vector: Optional[List[float]] = None,
-        limit: int = 15,
+        query_text: str | None = None,
+        query_vector: list[float] | None = None,
+        limit: int | None = 15,
         with_vector: bool = False,
-    ) -> List[ScoredResult]:
-        from qdrant_client.http.exceptions import UnexpectedResponse
-
+    ) -> list[ScoredResult]:
         if query_text is None and query_vector is None:
             raise MissingQueryParameterError()
 
@@ -176,47 +190,63 @@ class QDrantAdapter(VectorDBInterface):
         if query_vector is None:
             query_vector = (await self.embed_data([query_text]))[0]
 
+        client = None
         try:
             client = self.get_qdrant_client()
-            if limit == 0:
+            if limit is None:
                 collection_size = await client.count(collection_name=collection_name)
                 limit = collection_size.count
             if limit == 0:
+                await client.close()
                 return []
 
-            results = await client.search(
+            # Use query_points instead of search (API change in qdrant-client)
+            # query_points is the correct method for AsyncQdrantClient
+            query_result = await client.query_points(
                 collection_name=collection_name,
-                query_vector=models.NamedVector(
-                    name="text",
-                    vector=query_vector
-                    if query_vector is not None
-                    else (await self.embed_data([query_text]))[0],
+                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,
             )
 
             await client.close()
 
+            # Extract points from query_result
+            results = query_result.points
+
             return [
                 ScoredResult(
-                    id=parse_id(result.id),
+                    id=parse_id(str(result.id)),
                     payload={
                         **result.payload,
-                        "id": parse_id(result.id),
+                        "id": parse_id(str(result.id)),
                     },
-                    score=1 - result.score,
+                    score=1 - result.score if hasattr(result, "score") else 1.0,
                 )
                 for result in results
             ]
-        finally:
-            await client.close()
+        except Exception as e:
+            logger.error(f"Error in Qdrant search: {e}", exc_info=True)
+            if client:
+                await client.close()
+            return []
 
     async def batch_search(
         self,
         collection_name: str,
-        query_texts: List[str],
-        limit: int = None,
+        query_texts: list[str],
+        limit: int | None = None,
         with_vectors: bool = False,
     ):
         """
@@ -226,37 +256,59 @@ class QDrantAdapter(VectorDBInterface):
         - collection_name (str): Name of the collection to search in.
         - query_texts (List[str]): List of query texts to search for.
         - limit (int): List of result limits for search requests.
-        - with_vectors (bool, optional): Bool indicating whether to return vectors for search requests.
+        - with_vectors (bool, optional): Bool indicating whether to return
+                                         vectors for search requests.
 
         Returns:
         - results: The search results from Qdrant.
         """
 
-        vectors = await self.embed_data(query_texts)
-
-        # Generate dynamic search requests based on the provided embeddings
-        requests = [
-            models.SearchRequest(
-                vector=models.NamedVector(name="text", vector=vector),
-                limit=limit,
-                with_vector=with_vectors,
-            )
-            for vector in vectors
-        ]
+        client = self.get_qdrant_client()
+        if limit is None:
+            collection_size = await client.count(collection_name=collection_name)
+            limit = collection_size.count
+        if limit == 0:
+            await client.close()
+            return []
 
         client = self.get_qdrant_client()
 
-        # Perform batch search with the dynamically generated requests
-        results = await client.search_batch(
-            collection_name=collection_name, requests=requests
-        )
+        try:
+            # Use query_batch instead of search_batch (API change in qdrant-client)
+            # query_batch is the correct method for AsyncQdrantClient
+            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,
+            )
 
-        await client.close()
+            await client.close()
 
-        return [
-            filter(lambda result: result.score > 0.9, result_group)
-            for result_group in results
-        ]
+            # 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 []
+                filtered_points = [
+                    result for result in points if hasattr(result, "score") and result.score > 0.9
+                ]
+                filtered_results.append(filtered_points)
+
+            return filtered_results
+        except Exception as e:
+            logger.error(f"Error in Qdrant batch_search: {e}", exc_info=True)
+            await client.close()
+            return []
 
     async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
         client = self.get_qdrant_client()
@@ -269,6 +321,55 @@ 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()
+
+    async def get_collection_names(self) -> list[str]:
+        """
+        Get names of all collections in the database.
+
+        Returns:
+            list[str]: List of collection names.
+        """
+
+        client = self.get_qdrant_client()
+
+        response = await client.get_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()
+
+        return result

cognee_community_vector_adapter_qdrant-0.2.0/cognee_community_vector_adapter_qdrant/register.py

@@ -0,0 +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.0.3 → cognee_community_vector_adapter_qdrant-0.2.0}/pyproject.toml

@@ -1,10 +1,12 @@
 [project]
 name = "cognee-community-vector-adapter-qdrant"
-version = "0.0.3"
+version = "0.2.0"
 description = "Qdrant vector database adapter for cognee"
 readme = "README.md"
 requires-python = ">=3.11,<=3.13"
 dependencies = [
-    "qdrant-client>=1.14.2",
-    "cognee>=0.2.4",
+    "qdrant-client>=1.16.0",
+    "cognee==0.5.1",
+    "starlette>=0.48.0",
+    "instructor>=1.11"
 ]

cognee_community_vector_adapter_qdrant-0.0.3/cognee_community_vector_adapter_qdrant/register.py

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