cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py

@@ -0,0 +1,449 @@
+"""Neptune Analytics Hybrid Adapter combining Vector and Graph functionality"""
+
+import asyncio
+import json
+from typing import List, Optional, Any, Dict, Type, Tuple
+from uuid import UUID
+
+from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.databases.graph.neptune_driver.adapter import NeptuneGraphDB
+from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
+from cognee.infrastructure.engine import DataPoint
+from cognee.modules.storage.utils import JSONEncoder
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
+from cognee.infrastructure.databases.vector.models.PayloadSchema import PayloadSchema
+from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
+
+logger = get_logger("NeptuneAnalyticsAdapter")
+
+
+class IndexSchema(DataPoint):
+    """
+    Represents a schema for an index data point containing an ID and text.
+
+    Attributes:
+    - id: A string representing the unique identifier for the data point.
+    - text: A string representing the content of the data point.
+    - metadata: A dictionary with default index fields for the schema, currently configured
+    to include 'text'.
+    """
+
+    id: str
+    text: str
+    metadata: dict = {"index_fields": ["text"]}
+
+
+NEPTUNE_ANALYTICS_ENDPOINT_URL = "neptune-graph://"
+
+
+class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
+    """
+    Hybrid adapter that combines Neptune Analytics Vector and Graph functionality.
+
+    This adapter extends NeptuneGraphDB and implements VectorDBInterface to provide
+    a unified interface for working with Neptune Analytics as both a vector store
+    and a graph database.
+    """
+
+    _VECTOR_NODE_LABEL = "COGNEE_NODE"
+    _COLLECTION_PREFIX = "VECTOR_COLLECTION"
+    _TOPK_LOWER_BOUND = 0
+    _TOPK_UPPER_BOUND = 10
+
+    def __init__(
+        self,
+        graph_id: str,
+        embedding_engine: Optional[EmbeddingEngine] = None,
+        region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_session_token: Optional[str] = None,
+    ):
+        """
+        Initialize the Neptune Analytics hybrid adapter.
+
+        Parameters:
+        -----------
+            - graph_id (str): The Neptune Analytics graph identifier
+            - embedding_engine(Optional[EmbeddingEngine]): The embedding engine instance to translate text to vector.
+            - region (Optional[str]): AWS region where the graph is located (default: us-east-1)
+            - aws_access_key_id (Optional[str]): AWS access key ID
+            - aws_secret_access_key (Optional[str]): AWS secret access key
+            - aws_session_token (Optional[str]): AWS session token for temporary credentials
+        """
+        # Initialize the graph database functionality
+        super().__init__(
+            graph_id=graph_id,
+            region=region,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
+        )
+
+        # Add vector-specific attributes
+        self.embedding_engine = embedding_engine
+        logger.info(
+            f'Initialized Neptune Analytics hybrid adapter for graph: "{graph_id}" in region: "{self.region}"'
+        )
+
+    # VectorDBInterface methods implementation
+
+    async def get_connection(self):
+        """
+        This method is part of the default implementation but not defined in the interface.
+        No operation is performed and None will be returned here,
+        because the concept of connection is not applicable in this context.
+        """
+        return None
+
+    async def embed_data(self, data: list[str]) -> list[list[float]]:
+        """
+        Embeds the provided textual data into vector representation.
+
+        Uses the embedding engine to convert the list of strings into a list of float vectors.
+
+        Parameters:
+        -----------
+            - data (list[str]): A list of strings representing the data to be embedded.
+
+        Returns:
+        --------
+            - list[list[float]]: A list of embedded vectors corresponding to the input data.
+        """
+        self._validate_embedding_engine()
+        return await self.embedding_engine.embed_text(data)
+
+    async def has_collection(self, collection_name: str) -> bool:
+        """
+        Neptune Analytics stores vector on a node level,
+        so create_collection() implements interface for compliance but performs no operations when called.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the collection to check for existence.
+        Returns:
+        --------
+            - bool: Always return True.
+        """
+        return True
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        payload_schema: Optional[PayloadSchema] = None,
+    ):
+        """
+        Neptune Analytics stores vector on a node level, so create_collection() implements interface for compliance but performs no operations when called.
+        As the result, create_collection() will be no-op.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the new collection to create.
+            - payload_schema (Optional[PayloadSchema]): An optional schema for the payloads
+              within this collection. (default None)
+        """
+        pass
+
+    async def get_collection(self, collection_name: str):
+        """
+        This method is part of the default implementation but not defined in the interface.
+        No operation is performed here because the concept of collection is not applicable in NeptuneAnalytics vector store.
+        """
+        return None
+
+    async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
+        """
+        Insert new data points into the specified collection, by first inserting the node itself on the graph,
+        then execute neptune.algo.vectors.upsert() to insert the corresponded embedding.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the collection where data points will be added.
+            - data_points (List[DataPoint]): A list of data points to be added to the
+              collection.
+        """
+        self._validate_embedding_engine()
+
+        # Fetch embeddings
+        texts = [DataPoint.get_embeddable_data(t) for t in data_points]
+        data_vectors = await self.embedding_engine.embed_text(texts)
+
+        for index, data_point in enumerate(data_points):
+            node_id = data_point.id
+            # Fetch embedding from list instead
+            data_vector = data_vectors[index]
+
+            # Fetch properties
+            properties = self._serialize_properties(data_point.model_dump())
+            properties[self._COLLECTION_PREFIX] = collection_name
+            params = dict(
+                node_id=str(node_id),
+                properties=properties,
+                embedding=data_vector,
+                collection_name=collection_name,
+            )
+
+            # Compose the query and send
+            query_string = (
+                f"MERGE (n "
+                f":{self._VECTOR_NODE_LABEL} "
+                f" {{`~id`: $node_id}}) "
+                f"ON CREATE SET n = $properties, n.updated_at = timestamp() "
+                f"ON MATCH SET n += $properties, n.updated_at = timestamp() "
+                f"WITH n, $embedding AS embedding "
+                f"CALL neptune.algo.vectors.upsert(n, embedding) "
+                f"YIELD success "
+                f"RETURN success "
+            )
+
+            try:
+                self._client.query(query_string, params)
+            except Exception as e:
+                self._na_exception_handler(e, query_string)
+        pass
+
+    async def retrieve(self, collection_name: str, data_point_ids: list[str]):
+        """
+        Retrieve data points from a collection using their IDs.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the collection from which to retrieve data
+              points.
+            - data_point_ids (list[str]): A list of IDs of the data points to retrieve.
+        """
+        # Do the fetch for each node
+        params = dict(node_ids=data_point_ids, collection_name=collection_name)
+        query_string = (
+            f"MATCH( n :{self._VECTOR_NODE_LABEL}) "
+            f"WHERE id(n) in $node_ids AND "
+            f"n.{self._COLLECTION_PREFIX} = $collection_name "
+            f"RETURN n as payload "
+        )
+
+        try:
+            result = self._client.query(query_string, params)
+            return [self._get_scored_result(item) for item in result]
+        except Exception as e:
+            self._na_exception_handler(e, query_string)
+
+    async def search(
+        self,
+        collection_name: str,
+        query_text: Optional[str] = None,
+        query_vector: Optional[List[float]] = None,
+        limit: int = None,
+        with_vector: bool = False,
+    ):
+        """
+        Perform a search in the specified collection using either a text query or a vector
+        query.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the collection in which to perform the search.
+            - query_text (Optional[str]): An optional text query to search for in the
+              collection.
+            - query_vector (Optional[List[float]]): An optional vector representation for
+              searching the collection.
+            - limit (int): The maximum number of results to return from the search.
+            - with_vector (bool): Whether to return the vector representations with search
+              results, this is not supported for Neptune Analytics backend at the moment.
+
+        Returns:
+        --------
+            A list of scored results that match the query.
+        """
+        self._validate_embedding_engine()
+
+        if with_vector:
+            logger.warning(
+                "with_vector=True will include embedding vectors in the result. "
+                "This may trigger a resource-intensive query and increase response time. "
+                "Use this option only when vector data is required."
+            )
+
+        # In the case of excessive limit, or zero / negative value, limit will be set to 10.
+        if not limit or limit <= self._TOPK_LOWER_BOUND or limit > self._TOPK_UPPER_BOUND:
+            logger.warning(
+                "Provided limit (%s) is invalid (zero, negative, or exceeds maximum). "
+                "Defaulting to limit=10.",
+                limit,
+            )
+            limit = self._TOPK_UPPER_BOUND
+
+        if query_vector and query_text:
+            raise InvalidValueError(
+                message="The search function accepts either text or embedding as input, but not both."
+            )
+        elif query_text is None and query_vector is None:
+            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
+        elif query_vector:
+            embedding = query_vector
+        else:
+            data_vectors = await self.embedding_engine.embed_text([query_text])
+            embedding = data_vectors[0]
+
+        # Compose the parameters map
+        params = dict(embedding=embedding, param_topk=limit)
+        # Compose the query
+        query_string = f"""
+        CALL neptune.algo.vectors.topKByEmbeddingWithFiltering({{
+                topK: {limit},
+                embedding: {embedding}, 
+                nodeFilter: {{ equals: {{property: '{self._COLLECTION_PREFIX}', value: '{collection_name}'}} }}
+              }}
+            )
+        YIELD node, score
+        """
+
+        if with_vector:
+            query_string += """
+        WITH node, score, id(node) as node_id 
+        MATCH (n)
+        WHERE id(n) = id(node)
+        CALL neptune.algo.vectors.get(n)
+        YIELD embedding
+        RETURN node as payload, score, embedding
+        """
+
+        else:
+            query_string += """
+        RETURN node as payload, score
+        """
+
+        try:
+            query_response = self._client.query(query_string, params)
+            return [self._get_scored_result(item=item, with_score=True) for item in query_response]
+        except Exception as e:
+            self._na_exception_handler(e, query_string)
+
+    async def batch_search(
+        self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False
+    ):
+        """
+        Perform a batch search using multiple text queries against a collection.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the collection to conduct the batch search in.
+            - query_texts (List[str]): A list of text queries to use for the search.
+            - limit (int): The maximum number of results to return for each query.
+            - with_vectors (bool): Whether to include vector representations with search
+              results. (default False)
+
+        Returns:
+        --------
+            A list of search result sets, one for each query input.
+        """
+        self._validate_embedding_engine()
+
+        # Convert text to embedding array in batch
+        data_vectors = await self.embedding_engine.embed_text(query_texts)
+        return await asyncio.gather(
+            *[
+                self.search(collection_name, None, vector, limit, with_vectors)
+                for vector in data_vectors
+            ]
+        )
+
+    async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
+        """
+        Delete specified data points from a collection, by executing an OpenCypher query,
+        with matching [vector_label, collection_label, node_id] combination.
+
+        Parameters:
+        -----------
+            - collection_name (str): The name of the collection from which to delete data
+              points.
+            - data_point_ids (list[str]): A list of IDs of the data points to delete.
+        """
+        params = dict(node_ids=data_point_ids, collection_name=collection_name)
+        query_string = (
+            f"MATCH (n :{self._VECTOR_NODE_LABEL}) "
+            f"WHERE id(n) IN $node_ids "
+            f"AND n.{self._COLLECTION_PREFIX} = $collection_name "
+            f"DETACH DELETE n"
+        )
+        try:
+            self._client.query(query_string, params)
+        except Exception as e:
+            self._na_exception_handler(e, query_string)
+        pass
+
+    async def create_vector_index(self, index_name: str, index_property_name: str):
+        """
+        Neptune Analytics stores vectors at the node level,
+        so create_vector_index() implements the interface for compliance but performs no operation when called.
+        As a result, create_vector_index() invokes create_collection(), which is also a no-op.
+        This ensures the logic flow remains consistent, even if the concept of collections is introduced in a future release.
+        """
+        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]
+    ):
+        """
+        Indexes a list of data points into Neptune Analytics by creating them as nodes.
+
+        This method constructs a unique collection name by combining the `index_name` and
+        `index_property_name`, then delegates to `create_data_points()` to store the data.
+
+        Args:
+            index_name (str): The base name of the index.
+            index_property_name (str): The property name to append to the index name for uniqueness.
+            data_points (list[DataPoint]): A list of `DataPoint` instances to be indexed.
+
+        Returns:
+            None
+        """
+        await self.create_data_points(
+            f"{index_name}_{index_property_name}",
+            [
+                IndexSchema(
+                    id=str(data_point.id),
+                    text=getattr(data_point, data_point.metadata["index_fields"][0]),
+                )
+                for data_point in data_points
+            ],
+        )
+
+    async def prune(self):
+        """
+        Remove obsolete or unnecessary data from the database.
+        """
+        # Run actual truncate
+        self._client.query(f"MATCH (n :{self._VECTOR_NODE_LABEL}) DETACH DELETE n")
+        pass
+
+    @staticmethod
+    def _get_scored_result(
+        item: dict, with_vector: bool = False, with_score: bool = False
+    ) -> ScoredResult:
+        """
+        Util method to simplify the object creation of ScoredResult base on incoming NX payload response.
+        """
+        return ScoredResult(
+            id=item.get("payload").get("~id"),
+            payload=item.get("payload").get("~properties"),
+            score=item.get("score") if with_score else 0,
+            vector=item.get("embedding") if with_vector else None,
+        )
+
+    def _na_exception_handler(self, ex, query_string: str):
+        """
+        Generic exception handler for NA langchain.
+        """
+        logger.error("Neptune Analytics query failed: %s | Query: [%s]", ex, query_string)
+        raise ex
+
+    def _validate_embedding_engine(self):
+        """
+        Validates if the embedding_engine is defined
+        :raises: ValueError if this object does not have a valid embedding_engine
+        """
+        if self.embedding_engine is None:
+            raise ValueError(
+                "Neptune Analytics requires an embedder defined to make vector operations"
+            )

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -49,9 +49,17 @@ class SQLAlchemyAdapter:
 
                 run_sync(self.pull_from_s3())
 
-        self.engine = create_async_engine(
-            connection_string, poolclass=NullPool if "sqlite" in connection_string else None
-        )
+        if "sqlite" in connection_string:
+            self.engine = create_async_engine(
+                connection_string,
+                poolclass=NullPool,
+                connect_args={"timeout": 30},
+            )
+        else:
+            self.engine = create_async_engine(
+                connection_string, pool_size=12, max_overflow=12, poolclass=None
+            )
+
         self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)
 
     async def push_to_s3(self) -> None:

cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py

@@ -1,4 +1,5 @@
 import json
+import asyncio
 from uuid import UUID
 from typing import List, Optional
 from chromadb import AsyncHttpClient, Settings
@@ -161,6 +162,7 @@ class ChromaDBAdapter(VectorDBInterface):
         self.embedding_engine = embedding_engine
         self.url = url
         self.api_key = api_key
+        self.VECTOR_DB_LOCK = asyncio.Lock()
 
     async def get_connection(self) -> AsyncHttpClient:
         """
@@ -224,10 +226,13 @@ class ChromaDBAdapter(VectorDBInterface):
             - collection_name (str): The name of the collection to create.
             - payload_schema: The schema for the payload; can be None. (default None)
         """
-        client = await self.get_connection()
+        async with self.VECTOR_DB_LOCK:
+            client = await self.get_connection()
 
-        if not await self.has_collection(collection_name):
-            await client.create_collection(name=collection_name, metadata={"hnsw:space": "cosine"})
+            if not await self.has_collection(collection_name):
+                await client.create_collection(
+                    name=collection_name, metadata={"hnsw:space": "cosine"}
+                )
 
     async def get_collection(self, collection_name: str) -> AsyncHttpClient:
         """

cognee/infrastructure/databases/vector/create_vector_engine.py

@@ -19,7 +19,7 @@ def create_vector_engine(
     for each provider, raising an EnvironmentError if any are missing, or ImportError if the
     ChromaDB package is not installed.
 
-    Supported providers include: Weaviate, Qdrant, pgvector, FalkorDB, ChromaDB, and
+    Supported providers include: pgvector, FalkorDB, ChromaDB, and
     LanceDB.
 
     Parameters:
@@ -30,7 +30,7 @@ def create_vector_engine(
           providers.
         - vector_db_key (str): The API key or access token for the vector database instance.
         - vector_db_provider (str): The name of the vector database provider to use (e.g.,
-          'weaviate', 'qdrant').
+          'pgvector').
 
     Returns:
     --------
@@ -48,27 +48,7 @@ def create_vector_engine(
             embedding_engine=embedding_engine,
         )
 
-    if vector_db_provider == "weaviate":
-        from .weaviate_db import WeaviateAdapter
-
-        if not (vector_db_url and vector_db_key):
-            raise EnvironmentError("Missing requred Weaviate credentials!")
-
-        return WeaviateAdapter(vector_db_url, vector_db_key, embedding_engine=embedding_engine)
-
-    elif vector_db_provider == "qdrant":
-        if not (vector_db_url and vector_db_key):
-            raise EnvironmentError("Missing requred Qdrant credentials!")
-
-        from .qdrant.QDrantAdapter import QDrantAdapter
-
-        return QDrantAdapter(
-            url=vector_db_url,
-            api_key=vector_db_key,
-            embedding_engine=embedding_engine,
-        )
-
-    elif vector_db_provider == "pgvector":
+    if vector_db_provider == "pgvector":
         from cognee.infrastructure.databases.relational import get_relational_config
 
         # Get configuration for postgres database
@@ -122,6 +102,34 @@ def create_vector_engine(
             embedding_engine=embedding_engine,
         )
 
+    elif vector_db_provider == "neptune_analytics":
+        try:
+            from langchain_aws import NeptuneAnalyticsGraph
+        except ImportError:
+            raise ImportError(
+                "langchain_aws is not installed. Please install it with 'pip install langchain_aws'"
+            )
+
+        if not vector_db_url:
+            raise EnvironmentError("Missing Neptune endpoint.")
+
+        from cognee.infrastructure.databases.hybrid.neptune_analytics.NeptuneAnalyticsAdapter import (
+            NeptuneAnalyticsAdapter,
+            NEPTUNE_ANALYTICS_ENDPOINT_URL,
+        )
+
+        if not vector_db_url.startswith(NEPTUNE_ANALYTICS_ENDPOINT_URL):
+            raise ValueError(
+                f"Neptune endpoint must have the format '{NEPTUNE_ANALYTICS_ENDPOINT_URL}<GRAPH_ID>'"
+            )
+
+        graph_identifier = vector_db_url.replace(NEPTUNE_ANALYTICS_ENDPOINT_URL, "")
+
+        return NeptuneAnalyticsAdapter(
+            graph_id=graph_identifier,
+            embedding_engine=embedding_engine,
+        )
+
     else:
         from .lancedb.LanceDBAdapter import LanceDBAdapter
 

cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py

@@ -5,7 +5,9 @@ import litellm
 import os
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
-from cognee.infrastructure.llm.tokenizer.TikToken import TikTokenTokenizer
+from cognee.infrastructure.llm.tokenizer.TikToken import (
+    TikTokenTokenizer,
+)
 
 litellm.set_verbose = False
 logger = get_logger("FastembedEmbeddingEngine")

cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py

@@ -7,11 +7,19 @@ import litellm
 import os
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
-from cognee.infrastructure.llm.tokenizer.Gemini import GeminiTokenizer
-from cognee.infrastructure.llm.tokenizer.HuggingFace import HuggingFaceTokenizer
-from cognee.infrastructure.llm.tokenizer.Mistral import MistralTokenizer
-from cognee.infrastructure.llm.tokenizer.TikToken import TikTokenTokenizer
-from cognee.infrastructure.llm.embedding_rate_limiter import (
+from cognee.infrastructure.llm.tokenizer.Gemini import (
+    GeminiTokenizer,
+)
+from cognee.infrastructure.llm.tokenizer.HuggingFace import (
+    HuggingFaceTokenizer,
+)
+from cognee.infrastructure.llm.tokenizer.Mistral import (
+    MistralTokenizer,
+)
+from cognee.infrastructure.llm.tokenizer.TikToken import (
+    TikTokenTokenizer,
+)
+from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
     embedding_rate_limit_async,
     embedding_sleep_and_retry_async,
 )
@@ -177,7 +185,14 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
         elif "mistral" in self.provider.lower():
             tokenizer = MistralTokenizer(model=model, max_tokens=self.max_tokens)
         else:
-            tokenizer = HuggingFaceTokenizer(model=self.model, max_tokens=self.max_tokens)
+            try:
+                tokenizer = HuggingFaceTokenizer(
+                    model=self.model.replace("hosted_vllm/", ""), max_tokens=self.max_tokens
+                )
+            except Exception as e:
+                logger.warning(f"Could not get tokenizer from HuggingFace due to: {e}")
+                logger.info("Switching to TikToken default tokenizer.")
+                tokenizer = TikTokenTokenizer(model=None, max_tokens=self.max_tokens)
 
         logger.debug(f"Tokenizer loaded for model: {self.model}")
         return tokenizer

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -7,9 +7,10 @@ import os
 import aiohttp.http_exceptions
 
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
-from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
-from cognee.infrastructure.llm.tokenizer.HuggingFace import HuggingFaceTokenizer
-from cognee.infrastructure.llm.embedding_rate_limiter import (
+from cognee.infrastructure.llm.tokenizer.HuggingFace import (
+    HuggingFaceTokenizer,
+)
+from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
     embedding_rate_limit_async,
     embedding_sleep_and_retry_async,
 )

cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py

@@ -1,5 +1,7 @@
 from cognee.infrastructure.databases.vector.embeddings.config import get_embedding_config
-from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
 from .EmbeddingEngine import EmbeddingEngine
 from functools import lru_cache