cognee-community-vector-adapter-weaviate 0.0.2__tar.gz → 0.1.0__tar.gz

{cognee_community_vector_adapter_weaviate-0.0.2 → cognee_community_vector_adapter_weaviate-0.1.0}/PKG-INFO

@@ -1,13 +1,15 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: cognee-community-vector-adapter-weaviate
-Version: 0.0.2
+Version: 0.1.0
 Summary: Weaviate 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: cognee (==0.5.2)
+Requires-Dist: instructor (>=1.11)
+Requires-Dist: starlette (>=0.48.0)
 Requires-Dist: weaviate-client (>=4.9.6,<5.0.0)
 Description-Content-Type: text/markdown
 
@@ -55,7 +57,10 @@ self.client = weaviate.use_async_with_local(
 
 You can use the docker command provided by Weaviate (https://docs.weaviate.io/deploy/installation-guides/docker-installation)
 to run Weaviate with default settings. The command looks something like this, specifying the ports for connection:
-`docker run -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:1.32.4`
+
+```bash
+docker run -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:1.32.4
+```
 
 ## Usage
 

{cognee_community_vector_adapter_weaviate-0.0.2 → cognee_community_vector_adapter_weaviate-0.1.0}/README.md

@@ -42,7 +42,10 @@ self.client = weaviate.use_async_with_local(
 
 You can use the docker command provided by Weaviate (https://docs.weaviate.io/deploy/installation-guides/docker-installation)
 to run Weaviate with default settings. The command looks something like this, specifying the ports for connection:
-`docker run -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:1.32.4`
+
+```bash
+docker run -p 8080:8080 -p 50051:50051 cr.weaviate.io/semitechnologies/weaviate:1.32.4
+```
 
 ## Usage
 

{cognee_community_vector_adapter_weaviate-0.0.2 → cognee_community_vector_adapter_weaviate-0.1.0}/cognee_community_vector_adapter_weaviate/weaviate_adapter.py

@@ -1,26 +1,25 @@
 import asyncio
-from typing import List, Optional
 
-from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
-
-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.exceptions import CollectionNotFoundError
 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.databases.vector.vector_db_interface import VectorDBInterface
+from cognee.infrastructure.engine import DataPoint
+from cognee.infrastructure.engine.utils import parse_id
+from cognee.shared.logging_utils import get_logger
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
 
 logger = get_logger("WeaviateAdapter")
 
 
 def is_retryable_request(error):
-    from weaviate.exceptions import UnexpectedStatusCodeException
     from requests.exceptions import RequestException
 
+    from weaviate.exceptions import UnexpectedStatusCodeException
+
     if isinstance(error, UnexpectedStatusCodeException):
         # Retry on conflict, service unavailable, internal error
         return error.status_code in {409, 503, 500}
@@ -72,12 +71,19 @@ class WeaviateAdapter(VectorDBInterface):
     api_key: str
     embedding_engine: EmbeddingEngine = None
 
-    def __init__(self, url: str, api_key: str, embedding_engine: EmbeddingEngine):
+    def __init__(
+        self,
+        url: str,
+        api_key: str,
+        embedding_engine: EmbeddingEngine,
+        database_name: str = "cognee",
+    ):
         import weaviate
         import weaviate.classes as wvc
 
         self.url = url
         self.api_key = api_key
+        self.database_name = database_name
 
         self.embedding_engine = embedding_engine
         self.VECTOR_DB_LOCK = asyncio.Lock()
@@ -85,9 +91,7 @@ class WeaviateAdapter(VectorDBInterface):
         self.client = weaviate.use_async_with_weaviate_cloud(
             cluster_url=url,
             auth_credentials=weaviate.auth.AuthApiKey(api_key),
-            additional_config=wvc.init.AdditionalConfig(
-                timeout=wvc.init.Timeout(init=30)
-            ),
+            additional_config=wvc.init.AdditionalConfig(timeout=wvc.init.Timeout(init=30)),
         )
 
     async def get_client(self):
@@ -105,7 +109,7 @@ class WeaviateAdapter(VectorDBInterface):
 
         return self.client
 
-    async def embed_data(self, data: List[str]) -> List[float]:
+    async def embed_data(self, data: list[str]) -> list[float]:
         """
         Embed the given text data into vector representations.
 
@@ -187,7 +191,7 @@ class WeaviateAdapter(VectorDBInterface):
                 )
             else:
                 result = await self.get_collection(collection_name)
-                await client.close()
+                # await client.close()
                 return result
 
     async def get_collection(self, collection_name: str):
@@ -216,9 +220,7 @@ class WeaviateAdapter(VectorDBInterface):
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=2, min=1, max=6),
     )
-    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]):
         """
         Create or update data points in the specified collection in the Weaviate database.
 
@@ -269,9 +271,7 @@ class WeaviateAdapter(VectorDBInterface):
 
             return DataObject(uuid=data_point.id, properties=properties, vector=vector)
 
-        data_points = [
-            convert_to_weaviate_data_points(data_point) for data_point in data_points
-        ]
+        data_points = [convert_to_weaviate_data_points(data_point) for data_point in data_points]
 
         await self.get_client()
         collection = await self.get_collection(collection_name)
@@ -298,8 +298,8 @@ class WeaviateAdapter(VectorDBInterface):
         except Exception as error:
             logger.error("Error creating data points: %s", str(error))
             raise error
-        finally:
-            await self.client.close()
+        # finally:
+        #     await self.client.close()
 
     async def create_vector_index(self, index_name: str, index_property_name: str):
         """
@@ -381,22 +381,23 @@ class WeaviateAdapter(VectorDBInterface):
             data_point.id = data_point.uuid
             del data_point.properties
 
-        await self.client.close()
+        # await self.client.close()
         return data_points.objects
 
     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,
+        include_payload: bool = False,
     ):
         """
         Perform a search on a collection using either a text query or a vector query.
 
-        Return scored results based on the search criteria provided. Raise MissingQueryParameterError if
-        no query is provided.
+        Return scored results based on the search criteria provided.
+        Raise MissingQueryParameterError if no query is provided.
 
         Parameters:
         -----------
@@ -408,6 +409,7 @@ class WeaviateAdapter(VectorDBInterface):
               searching. (default None)
             - limit (int): The maximum number of results to return. (default 15)
             - with_vector (bool): Include vector information in the results. (default False)
+            - include_payload (bool): Include payload information in the results. (default False)
 
         Returns:
         --------
@@ -423,28 +425,32 @@ class WeaviateAdapter(VectorDBInterface):
         if query_vector is None:
             query_vector = (await self.embed_data([query_text]))[0]
 
-        # TODO: Creation of new client for every search call. This is VERY ugly, needs discussion. (Andrej's comment)
+        # TODO: Creation of new client for every search call. This is VERY ugly. Should change.
         async with weaviate.use_async_with_weaviate_cloud(
             cluster_url=self.url,
             auth_credentials=weaviate.auth.AuthApiKey(self.api_key),
-            additional_config=wvc.init.AdditionalConfig(
-                timeout=wvc.init.Timeout(init=30)
-            ),
+            additional_config=wvc.init.AdditionalConfig(timeout=wvc.init.Timeout(init=30)),
         ) as client:
             if not await client.collections.exists(collection_name):
-                raise CollectionNotFoundError(
-                    f"Collection '{collection_name}' not found."
-                )
+                raise CollectionNotFoundError(f"Collection '{collection_name}' not found.")
 
             collection = client.collections.get(collection_name)
 
+            if limit is None:
+                result = await collection.aggregate.over_all(total_count=True)
+                limit = result.total_count
+
+            if limit == 0:
+                return []
+
             try:
                 search_result = await collection.query.hybrid(
                     query=None,
                     vector=query_vector,
-                    limit=limit if limit > 0 else None,
+                    limit=limit,
                     include_vector=with_vector,
                     return_metadata=wvc.query.MetadataQuery(score=True),
+                    return_properties=include_payload,
                 )
 
                 return [
@@ -462,9 +468,10 @@ class WeaviateAdapter(VectorDBInterface):
     async def batch_search(
         self,
         collection_name: str,
-        query_texts: List[str],
-        limit: int,
+        query_texts: list[str],
+        limit: int | None,
         with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         """
         Execute a batch search for multiple query texts in the specified collection.
@@ -479,6 +486,7 @@ class WeaviateAdapter(VectorDBInterface):
             - limit (int): The maximum number of results to return for each query.
             - with_vectors (bool): Indicate whether to include vector information in the
               results. (default False)
+            - include_payload (bool): Include payload information in the results.
 
         Returns:
         --------
@@ -508,11 +516,11 @@ class WeaviateAdapter(VectorDBInterface):
                 query_vector=query_vector,
                 limit=limit,
                 with_vector=with_vectors,
+                include_payload=include_payload,
             )
 
         return [
-            await query_search(query_vector)
-            for query_vector in await self.embed_data(query_texts)
+            await query_search(query_vector) for query_vector in await self.embed_data(query_texts)
         ]
 
     async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
@@ -542,7 +550,7 @@ class WeaviateAdapter(VectorDBInterface):
             filters=Filter.by_id().contains_any(data_point_ids)
         )
 
-        await self.client.close()
+        # await self.client.close()
         return result
 
     async def prune(self):
@@ -553,4 +561,15 @@ class WeaviateAdapter(VectorDBInterface):
         """
         client = await self.get_client()
         await client.collections.delete_all()
-        await client.close()
+        # 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 = await self.get_client()
+        return await client.collections.list_all()

{cognee_community_vector_adapter_weaviate-0.0.2 → cognee_community_vector_adapter_weaviate-0.1.0}/pyproject.toml

@@ -1,10 +1,12 @@
 [project]
 name = "cognee-community-vector-adapter-weaviate"
-version = "0.0.2"
+version = "0.1.0"
 description = "Weaviate vector database adapter for cognee"
 readme = "README.md"
 requires-python = ">=3.11,<=3.13"
 dependencies = [
     "weaviate-client>=4.9.6,<5.0.0",
-    "cognee>=0.2.4"
+    "cognee==0.5.2",
+    "starlette>=0.48.0",
+    "instructor>=1.11"
 ]