cognee-community-vector-adapter-weaviate 0.0.2__py3-none-any.whl

cognee_community_vector_adapter_weaviate/__init__.py

@@ -0,0 +1,3 @@
+from .weaviate_adapter import WeaviateAdapter
+
+__all__ = ["WeaviateAdapter"]

cognee_community_vector_adapter_weaviate/register.py

@@ -0,0 +1,5 @@
+from cognee.infrastructure.databases.vector import use_vector_adapter
+
+from .weaviate_adapter import WeaviateAdapter
+
+use_vector_adapter("weaviate", WeaviateAdapter)

cognee_community_vector_adapter_weaviate/weaviate_adapter.py

@@ -0,0 +1,556 @@
+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.models.ScoredResult import ScoredResult
+from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
+
+logger = get_logger("WeaviateAdapter")
+
+
+def is_retryable_request(error):
+    from weaviate.exceptions import UnexpectedStatusCodeException
+    from requests.exceptions import RequestException
+
+    if isinstance(error, UnexpectedStatusCodeException):
+        # Retry on conflict, service unavailable, internal error
+        return error.status_code in {409, 503, 500}
+    if isinstance(error, RequestException):
+        return True  # Includes timeout, connection error, etc.
+    return False
+
+
+class IndexSchema(DataPoint):
+    """
+    Define a schema for indexing data points with textual content.
+
+    The IndexSchema class inherits from DataPoint and includes the following public
+    attributes:
+
+    - text: A string representing the main content of the data point.
+    - metadata: A dictionary containing indexing information, specifically the fields to be
+    indexed (in this case, the 'text' field).
+    """
+
+    text: str
+
+    metadata: dict = {"index_fields": ["text"]}
+
+
+class WeaviateAdapter(VectorDBInterface):
+    """
+    Adapt the Weaviate vector database to an interface for managing collections and data
+    points.
+
+    Public methods:
+    - get_client
+    - embed_data
+    - has_collection
+    - create_collection
+    - get_collection
+    - create_data_points
+    - create_vector_index
+    - index_data_points
+    - retrieve
+    - search
+    - batch_search
+    - delete_data_points
+    - prune
+    """
+
+    name = "Weaviate"
+    url: str
+    api_key: str
+    embedding_engine: EmbeddingEngine = None
+
+    def __init__(self, url: str, api_key: str, embedding_engine: EmbeddingEngine):
+        import weaviate
+        import weaviate.classes as wvc
+
+        self.url = url
+        self.api_key = api_key
+
+        self.embedding_engine = embedding_engine
+        self.VECTOR_DB_LOCK = asyncio.Lock()
+
+        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)
+            ),
+        )
+
+    async def get_client(self):
+        """
+        Establish a connection to the Weaviate client.
+
+        Return the Weaviate client instance after connecting asynchronously.
+
+        Returns:
+        --------
+
+            The Weaviate client instance.
+        """
+        await self.client.connect()
+
+        return self.client
+
+    async def embed_data(self, data: List[str]) -> List[float]:
+        """
+        Embed the given text data into vector representations.
+
+        Given a list of strings, return their vector embeddings using the configured embedding
+        engine.
+
+        Parameters:
+        -----------
+
+            - data (List[str]): A list of strings to be embedded.
+
+        Returns:
+        --------
+
+            - List[float]: A list of float vectors corresponding to the embedded text data.
+        """
+        return await self.embedding_engine.embed_text(data)
+
+    async def has_collection(self, collection_name: str) -> bool:
+        """
+        Check if a collection exists in the Weaviate database.
+
+        Return a boolean indicating the presence of the specified collection.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection to check.
+
+        Returns:
+        --------
+
+            - bool: True if the collection exists, otherwise False.
+        """
+        return await self.client.collections.exists(collection_name)
+
+    @retry(
+        retry=retry_if_exception(is_retryable_request),
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=2, min=1, max=6),
+    )
+    async def create_collection(
+        self,
+        collection_name: str,
+        payload_schema=None,
+    ):
+        """
+        Create a new collection in the Weaviate database if it does not already exist.
+
+        The collection will be initialized with a default schema.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the new collection to be created.
+            - payload_schema: Optional schema definition for the collection payload. (default
+              None)
+
+        Returns:
+        --------
+
+            The created collection's configuration, if a new collection was made, otherwise
+            information about the existing collection.
+        """
+        import weaviate.classes.config as wvcc
+
+        client = await self.get_client()
+        async with self.VECTOR_DB_LOCK:
+            if not await self.has_collection(collection_name):
+                return await client.collections.create(
+                    name=collection_name,
+                    properties=[
+                        wvcc.Property(
+                            name="text",
+                            data_type=wvcc.DataType.TEXT,
+                            skip_vectorization=True,
+                        )
+                    ],
+                )
+            else:
+                result = await self.get_collection(collection_name)
+                await client.close()
+                return result
+
+    async def get_collection(self, collection_name: str):
+        """
+        Retrieve a collection from the Weaviate database by its name.
+
+        Raise a CollectionNotFoundError if the specified collection does not exist.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection to be retrieved.
+
+        Returns:
+        --------
+
+            The requested collection object from the database.
+        """
+        if not await self.has_collection(collection_name):
+            raise CollectionNotFoundError(f"Collection '{collection_name}' not found.")
+
+        return self.client.collections.get(collection_name)
+
+    @retry(
+        retry=retry_if_exception(is_retryable_request),
+        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]
+    ):
+        """
+        Create or update data points in the specified collection in the Weaviate database.
+
+        Process the list of data points, embedding them and either inserting them or updating if
+        they already exist.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection to add data points to.
+            - data_points (List[DataPoint]): A list of DataPoint objects to be created or
+              updated in the collection.
+
+        Returns:
+        --------
+
+            Information about the inserted or updated data points in the collection.
+        """
+        from weaviate.classes.data import DataObject
+
+        data_vectors = await self.embed_data(
+            [DataPoint.get_embeddable_data(data_point) for data_point in data_points]
+        )
+
+        def convert_to_weaviate_data_points(data_point: DataPoint):
+            """
+            Transform a DataPoint object into a Weaviate DataObject format for insertion.
+
+            Return a DataObject ready for use in Weaviate with the properties and vector included.
+
+            Parameters:
+            -----------
+
+                - data_point (DataPoint): The DataPoint to convert into the Weaviate DataObject
+                  format.
+
+            Returns:
+            --------
+
+                The corresponding Weaviate DataObject representing the data point.
+            """
+            vector = data_vectors[data_points.index(data_point)]
+            properties = data_point.model_dump()
+
+            if "id" in properties:
+                properties["uuid"] = str(data_point.id)
+                del properties["id"]
+
+            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
+        ]
+
+        await self.get_client()
+        collection = await self.get_collection(collection_name)
+
+        try:
+            if len(data_points) > 1:
+                return await collection.data.insert_many(data_points)
+            else:
+                data_point: DataObject = data_points[0]
+                if await collection.data.exists(data_point.uuid):
+                    return await collection.data.update(
+                        uuid=data_point.uuid,
+                        vector=data_point.vector,
+                        properties=data_point.properties,
+                        references=data_point.references,
+                    )
+                else:
+                    return await collection.data.insert(
+                        uuid=data_point.uuid,
+                        vector=data_point.vector,
+                        properties=data_point.properties,
+                        references=data_point.references,
+                    )
+        except Exception as error:
+            logger.error("Error creating data points: %s", str(error))
+            raise error
+        finally:
+            await self.client.close()
+
+    async def create_vector_index(self, index_name: str, index_property_name: str):
+        """
+        Create a vector index based on an index name and property name by creating a
+        corresponding collection.
+
+        Parameters:
+        -----------
+
+            - index_name (str): The name for the vector index.
+            - index_property_name (str): The property name associated with the vector index.
+
+        Returns:
+        --------
+
+            The created collection representing the vector index.
+        """
+        return await self.create_collection(f"{index_name}_{index_property_name}")
+
+    async def index_data_points(
+        self, index_name: str, index_property_name: str, data_points: list[DataPoint]
+    ):
+        """
+        Index a list of data points by creating an associated vector index collection.
+
+        Data points are transformed into embeddable data before being processed for indexing.
+
+        Parameters:
+        -----------
+
+            - index_name (str): The index name under which to store the data points.
+            - index_property_name (str): The associated property name for the index.
+            - data_points (list[DataPoint]): A list of DataPoint objects to be indexed.
+
+        Returns:
+        --------
+
+            Information about the operation of indexing the data points.
+        """
+        return await self.create_data_points(
+            f"{index_name}_{index_property_name}",
+            [
+                IndexSchema(
+                    id=data_point.id,
+                    text=DataPoint.get_embeddable_data(data_point),
+                )
+                for data_point in data_points
+            ],
+        )
+
+    async def retrieve(self, collection_name: str, data_point_ids: list[str]):
+        """
+        Fetch data points from a specified collection based on their IDs.
+
+        Return data points wrapped in an object containing their properties after
+        transformation.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection to retrieve data points from.
+            - data_point_ids (list[str]): A list of IDs for the data points to retrieve.
+
+        Returns:
+        --------
+
+            A list of objects representing the retrieved data points.
+        """
+        from weaviate.classes.query import Filter
+
+        await self.get_client()
+        collection = await self.get_collection(collection_name)
+        data_points = await collection.query.fetch_objects(
+            filters=Filter.by_id().contains_any(data_point_ids)
+        )
+
+        for data_point in data_points.objects:
+            data_point.payload = data_point.properties
+            data_point.id = data_point.uuid
+            del data_point.properties
+
+        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,
+        with_vector: 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.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection to search within.
+            - query_text (Optional[str]): Optional plain text query for searching. (default
+              None)
+            - query_vector (Optional[List[float]]): Optional vector representation for
+              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)
+
+        Returns:
+        --------
+
+            A list of scored results matching the search criteria.
+        """
+        import weaviate.classes as wvc
+        import weaviate.exceptions
+
+        if query_text is None and query_vector is None:
+            raise MissingQueryParameterError()
+
+        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)
+        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)
+            ),
+        ) as client:
+            if not await client.collections.exists(collection_name):
+                raise CollectionNotFoundError(
+                    f"Collection '{collection_name}' not found."
+                )
+
+            collection = client.collections.get(collection_name)
+
+            try:
+                search_result = await collection.query.hybrid(
+                    query=None,
+                    vector=query_vector,
+                    limit=limit if limit > 0 else None,
+                    include_vector=with_vector,
+                    return_metadata=wvc.query.MetadataQuery(score=True),
+                )
+
+                return [
+                    ScoredResult(
+                        id=parse_id(str(result.uuid)),
+                        payload=result.properties,
+                        score=1 - float(result.metadata.score),
+                    )
+                    for result in search_result.objects
+                ]
+            except weaviate.exceptions.WeaviateInvalidInputError:
+                # Ignore if the collection doesn't exist
+                return []
+
+    async def batch_search(
+        self,
+        collection_name: str,
+        query_texts: List[str],
+        limit: int,
+        with_vectors: bool = False,
+    ):
+        """
+        Execute a batch search for multiple query texts in the specified collection.
+
+        Return a list of results for each query performed in parallel.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection to search within.
+            - query_texts (List[str]): A list of text queries to be processed in a batch.
+            - 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)
+
+        Returns:
+        --------
+
+            A list containing results for each search query executed.
+        """
+
+        def query_search(query_vector):
+            """
+            Wrap the search operation based on a query vector for fetching results.
+
+            This function coordinates the search call, ensuring the collection name and search
+            parameters are applied.
+
+            Parameters:
+            -----------
+
+                - query_vector: The vector representation of the query for searching.
+
+            Returns:
+            --------
+
+                The results of the search operation on the specified collection.
+            """
+            return self.search(
+                collection_name,
+                query_vector=query_vector,
+                limit=limit,
+                with_vector=with_vectors,
+            )
+
+        return [
+            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]):
+        """
+        Remove specified data points from a collection based on their IDs.
+
+        Return information about the deletion result, ideally confirming the operation's
+        success.
+
+        Parameters:
+        -----------
+
+            - collection_name (str): The name of the collection from which to delete data
+              points.
+            - data_point_ids (list[str]): A list of IDs for the data points to be deleted.
+
+        Returns:
+        --------
+
+            Confirmation of deletion operation result.
+        """
+        from weaviate.classes.query import Filter
+
+        await self.get_client()
+        collection = await self.get_collection(collection_name)
+        result = await collection.data.delete_many(
+            filters=Filter.by_id().contains_any(data_point_ids)
+        )
+
+        await self.client.close()
+        return result
+
+    async def prune(self):
+        """
+        Delete all collections from the Weaviate database.
+
+        This operation will remove all data and cannot be undone.
+        """
+        client = await self.get_client()
+        await client.collections.delete_all()
+        await client.close()

cognee_community_vector_adapter_weaviate-0.0.2.dist-info/METADATA

@@ -0,0 +1,138 @@
+Metadata-Version: 2.3
+Name: cognee-community-vector-adapter-weaviate
+Version: 0.0.2
+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: weaviate-client (>=4.9.6,<5.0.0)
+Description-Content-Type: text/markdown
+
+# Cognee Community Weaviate Vector Adapter
+
+This is a community-maintained adapter that enables Cognee to work with Weaviate as a vector database.
+
+## Installation
+
+If published, the package can be simply installed via pip:
+
+```bash
+pip install cognee-community-vector-adapter-weaviate
+```
+
+In case it is not published yet, you can use poetry to locally build the adapter package:
+
+```bash
+pip install poetry
+poetry install # run this command in the directory containing the pyproject.toml file
+```
+
+## Connection Setup
+The provided code creates an async client connected to a remote instance of Weaviate. If you want to connect to a local 
+instance, like running a docker container locally and connecting to it, you need to change a few lines of code. 
+In the `weaviate_adapter.py` file inside the `.../weaviate/cognee_community_vector_adapter_weaviate` directory, replace 
+the following lines in the constructor: 
+
+```
+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)),
+)
+```
+with the following:
+
+```
+self.client = weaviate.use_async_with_local(
+    host="localhost",
+    port=8080,
+    grpc_port=50051
+)
+```
+
+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`
+
+## Usage
+
+```python
+import asyncio
+import os
+from cognee import config, prune, add, cognify, search, SearchType
+
+# Import the register module to enable Weaviate support
+import cognee_community_vector_adapter_weaviate.register
+
+async def main():
+    # Configure databases
+    config.set_relational_db_config({
+        "db_provider": "sqlite",
+    })
+    config.set_vector_db_config({
+        "vector_db_provider": "weaviate",
+        "vector_db_url": os.getenv("VECTOR_DB_URL"),  # or your Weaviate URL
+        "vector_db_key": os.getenv("VECTOR_DB_KEY"),  # or your API key
+    })
+    config.set_graph_db_config({
+        "graph_database_provider": "networkx",
+    })
+    
+    # Optional: Clean previous data
+    await prune.prune_data()
+    await prune.prune_system()
+    
+    # Add and process your content
+    text = "Your text content here"
+    await add(text)
+    await cognify()
+    
+    # Search
+    search_results = await search(
+        query_type=SearchType.GRAPH_COMPLETION,
+        query_text="Your search query"
+    )
+    
+    for result in search_results:
+        print(result)
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+## Configuration
+
+The Weaviate adapter requires the following configuration parameters:
+
+- `vector_db_url`: Your Weaviate cluster endpoint URL
+- `vector_db_key`: Your Weaviate API key
+- `vector_db_provider`: Set to "weaviate"
+
+### Environment Variables
+
+Set the following environment variables or pass them directly in the config:
+
+```bash
+export VECTOR_DB_URL="https://your-weaviate-instance.weaviate.network"
+export VECTOR_DB_KEY="your-api-key"
+```
+
+**Alternative:** You can also use the [`.env.template`](https://github.com/topoteretes/cognee/blob/main/.env.template) file from the main cognee repository. Copy it to your project directory, rename it to `.env`, and fill in your Weaviate configuration values.
+
+## Requirements
+
+- Python >= 3.11, <= 3.13
+- weaviate-client >= 4.9.6, < 5.0.0
+- cognee >= 0.2.1
+
+## Features
+
+- Full vector search capabilities
+- Batch operations support
+- Async/await support
+- Retry logic for better reliability
+- Collection management
+- Data point indexing and retrieval 

cognee_community_vector_adapter_weaviate-0.0.2.dist-info/RECORD

@@ -0,0 +1,6 @@
+cognee_community_vector_adapter_weaviate/__init__.py,sha256=EWEnTWaRf8kw8rNnbmppLYCiNl00Ltv9hsb1NmlIgMw,77
+cognee_community_vector_adapter_weaviate/register.py,sha256=aGTpOdpiFXYlGkDshE5XjLC_iJRxD66nSz66SnuaxpY,166
+cognee_community_vector_adapter_weaviate/weaviate_adapter.py,sha256=rDZBXQFUyU6sgVcbmInwK3JPs3oaa6Y4Ek4BzaLE9vI,18594
+cognee_community_vector_adapter_weaviate-0.0.2.dist-info/METADATA,sha256=6AvHIbq5KJTWtvmfvBvn38W4L8V10wj7zrDZJJD-hBg,4252
+cognee_community_vector_adapter_weaviate-0.0.2.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+cognee_community_vector_adapter_weaviate-0.0.2.dist-info/RECORD,,

cognee_community_vector_adapter_weaviate-0.0.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.2
+Root-Is-Purelib: true
+Tag: py3-none-any