cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2.dev0__py3-none-any.whl

cognee/infrastructure/databases/vector/create_vector_engine.py

@@ -7,7 +7,6 @@ from cognee.infrastructure.databases.graph.config import get_graph_context_confi
 from functools import lru_cache
 
 
-@lru_cache
 def create_vector_engine(
     vector_db_provider: str,
     vector_db_url: str,
@@ -15,6 +14,38 @@ def create_vector_engine(
     vector_db_port: str = "",
     vector_db_key: str = "",
     vector_dataset_database_handler: str = "",
+    vector_db_username: str = "",
+    vector_db_password: str = "",
+    vector_db_host: str = "",
+):
+    """
+    Wrapper function to call create vector engine with caching.
+    For a detailed description, see _create_vector_engine.
+    """
+    return _create_vector_engine(
+        vector_db_provider,
+        vector_db_url,
+        vector_db_name,
+        vector_db_port,
+        vector_db_key,
+        vector_dataset_database_handler,
+        vector_db_username,
+        vector_db_password,
+        vector_db_host,
+    )
+
+
+@lru_cache
+def _create_vector_engine(
+    vector_db_provider: str,
+    vector_db_url: str,
+    vector_db_name: str,
+    vector_db_port: str,
+    vector_db_key: str,
+    vector_dataset_database_handler: str,
+    vector_db_username: str,
+    vector_db_password: str,
+    vector_db_host: str,
 ):
     """
     Create a vector database engine based on the specified provider.
@@ -55,27 +86,43 @@ def create_vector_engine(
         )
 
     if vector_db_provider.lower() == "pgvector":
-        from cognee.infrastructure.databases.relational import get_relational_config
-
-        # Get configuration for postgres database
-        relational_config = get_relational_config()
-        db_username = relational_config.db_username
-        db_password = relational_config.db_password
-        db_host = relational_config.db_host
-        db_port = relational_config.db_port
-        db_name = relational_config.db_name
-
-        if not (db_host and db_port and db_name and db_username and db_password):
-            raise EnvironmentError("Missing requred pgvector credentials!")
-
-        connection_string = URL.create(
-            "postgresql+asyncpg",
-            username=db_username,
-            password=db_password,
-            host=db_host,
-            port=int(db_port),
-            database=db_name,
-        )
+        from cognee.context_global_variables import backend_access_control_enabled
+
+        if backend_access_control_enabled():
+            connection_string: str = (
+                f"postgresql+asyncpg://{vector_db_username}:{vector_db_password}"
+                f"@{vector_db_host}:{vector_db_port}/{vector_db_name}"
+            )
+        else:
+            if (
+                vector_db_port
+                and vector_db_username
+                and vector_db_password
+                and vector_db_host
+                and vector_db_name
+            ):
+                connection_string: str = (
+                    f"postgresql+asyncpg://{vector_db_username}:{vector_db_password}"
+                    f"@{vector_db_host}:{vector_db_port}/{vector_db_name}"
+                )
+            else:
+                from cognee.infrastructure.databases.relational import get_relational_config
+
+                # Get configuration for postgres database
+                relational_config = get_relational_config()
+                db_username = relational_config.db_username
+                db_password = relational_config.db_password
+                db_host = relational_config.db_host
+                db_port = relational_config.db_port
+                db_name = relational_config.db_name
+
+                if not (db_host and db_port and db_name and db_username and db_password):
+                    raise EnvironmentError("Missing required pgvector credentials!")
+
+                connection_string: str = (
+                    f"postgresql+asyncpg://{db_username}:{db_password}"
+                    f"@{db_host}:{db_port}/{db_name}"
+                )
 
         try:
             from .pgvector.PGVectorAdapter import PGVectorAdapter

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

@@ -14,6 +14,8 @@ from tenacity import (
 )
 import litellm
 import os
+from urllib.parse import urlparse
+import httpx
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.databases.exceptions import EmbeddingException
 from cognee.infrastructure.llm.tokenizer.HuggingFace import (
@@ -79,10 +81,26 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
             enable_mocking = str(enable_mocking).lower()
         self.mock = enable_mocking in ("true", "1", "yes")
 
+        # Validate provided custom embedding endpoint early to avoid long hangs later
+        if self.endpoint:
+            try:
+                parsed = urlparse(self.endpoint)
+            except Exception:
+                parsed = None
+            if not parsed or parsed.scheme not in ("http", "https") or not parsed.netloc:
+                logger.error(
+                    "Invalid EMBEDDING_ENDPOINT configured: '%s'. Expected a URL starting with http:// or https://",
+                    str(self.endpoint),
+                )
+                raise EmbeddingException(
+                    "Invalid EMBEDDING_ENDPOINT. Please set a valid URL (e.g., https://host:port) "
+                    "via environment variable EMBEDDING_ENDPOINT."
+                )
+
     @retry(
         stop=stop_after_delay(128),
         wait=wait_exponential_jitter(2, 128),
-        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        retry=retry_if_not_exception_type((litellm.exceptions.NotFoundError)),
         before_sleep=before_sleep_log(logger, logging.DEBUG),
         reraise=True,
     )
@@ -111,12 +129,21 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
                 return [data["embedding"] for data in response["data"]]
             else:
                 async with embedding_rate_limiter_context_manager():
-                    response = await litellm.aembedding(
-                        model=self.model,
-                        input=text,
-                        api_key=self.api_key,
-                        api_base=self.endpoint,
-                        api_version=self.api_version,
+                    embedding_kwargs = {
+                        "model": self.model,
+                        "input": text,
+                        "api_key": self.api_key,
+                        "api_base": self.endpoint,
+                        "api_version": self.api_version,
+                    }
+                    # Pass through target embedding dimensions when supported
+                    if self.dimensions is not None:
+                        embedding_kwargs["dimensions"] = self.dimensions
+
+                    # Ensure each attempt does not hang indefinitely
+                    response = await asyncio.wait_for(
+                        litellm.aembedding(**embedding_kwargs),
+                        timeout=30.0,
                     )
 
                 return [data["embedding"] for data in response.data]
@@ -154,6 +181,27 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
             logger.error("Context window exceeded for embedding text: %s", str(error))
             raise error
 
+        except asyncio.TimeoutError as e:
+            # Per-attempt timeout – likely an unreachable endpoint
+            logger.error(
+                "Embedding endpoint timed out. EMBEDDING_ENDPOINT='%s'. "
+                "Verify that the endpoint is reachable and correct.",
+                str(self.endpoint),
+            )
+            raise EmbeddingException(
+                "Embedding request timed out. Check EMBEDDING_ENDPOINT connectivity."
+            ) from e
+
+        except (httpx.ConnectError, httpx.ReadTimeout) as e:
+            logger.error(
+                "Failed to connect to embedding endpoint. EMBEDDING_ENDPOINT='%s'. "
+                "Ensure the URL is correct and the server is running.",
+                str(self.endpoint),
+            )
+            raise EmbeddingException(
+                "Cannot connect to embedding endpoint. Check EMBEDDING_ENDPOINT."
+            ) from e
+
         except (
             litellm.exceptions.BadRequestError,
             litellm.exceptions.NotFoundError,
@@ -162,8 +210,15 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
             raise EmbeddingException(f"Failed to index data points using model {self.model}") from e
 
         except Exception as error:
-            logger.error("Error embedding text: %s", str(error))
-            raise error
+            # Fall back to a clear, actionable message for connectivity/misconfiguration issues
+            logger.error(
+                "Error embedding text: %s. EMBEDDING_ENDPOINT='%s'.",
+                str(error),
+                str(self.endpoint),
+            )
+            raise EmbeddingException(
+                "Embedding failed due to an unexpected error. Verify EMBEDDING_ENDPOINT and provider settings."
+            ) from error
 
     def get_vector_size(self) -> int:
         """

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

@@ -57,7 +57,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         model: Optional[str] = "avr/sfr-embedding-mistral:latest",
         dimensions: Optional[int] = 1024,
         max_completion_tokens: int = 512,
-        endpoint: Optional[str] = "http://localhost:11434/api/embeddings",
+        endpoint: Optional[str] = "http://localhost:11434/api/embed",
         huggingface_tokenizer: str = "Salesforce/SFR-Embedding-Mistral",
         batch_size: int = 100,
     ):
@@ -93,6 +93,10 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         if self.mock:
             return [[0.0] * self.dimensions for _ in text]
 
+        # Handle case when a single string is passed instead of a list
+        if not isinstance(text, list):
+            text = [text]
+
         embeddings = await asyncio.gather(*[self._get_embedding(prompt) for prompt in text])
         return embeddings
 
@@ -107,7 +111,12 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         """
         Internal method to call the Ollama embeddings endpoint for a single prompt.
         """
-        payload = {"model": self.model, "prompt": prompt, "input": prompt}
+        payload = {
+            "model": self.model,
+            "prompt": prompt,
+            "input": prompt,
+            "dimensions": self.dimensions,
+        }
 
         headers = {}
         api_key = os.getenv("LLM_API_KEY")
@@ -124,6 +133,8 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
                     data = await response.json()
                     if "embeddings" in data:
                         return data["embeddings"][0]
+                    if "embedding" in data:
+                        return data["embedding"]
                     else:
                         return data["data"][0]["embedding"]
 

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -231,6 +231,7 @@ class LanceDBAdapter(VectorDBInterface):
         limit: Optional[int] = 15,
         with_vector: bool = False,
         normalized: bool = True,
+        include_payload: bool = False,
     ):
         if query_text is None and query_vector is None:
             raise MissingQueryParameterError()
@@ -247,17 +248,27 @@ class LanceDBAdapter(VectorDBInterface):
         if limit <= 0:
             return []
 
-        result_values = await collection.vector_search(query_vector).limit(limit).to_list()
+        # Note: Exclude payload if not needed to optimize performance
+        select_columns = (
+            ["id", "vector", "payload", "_distance"]
+            if include_payload
+            else ["id", "vector", "_distance"]
+        )
+        result_values = (
+            await collection.vector_search(query_vector)
+            .select(select_columns)
+            .limit(limit)
+            .to_list()
+        )
 
         if not result_values:
             return []
-
         normalized_values = normalize_distances(result_values)
 
         return [
             ScoredResult(
                 id=parse_id(result["id"]),
-                payload=result["payload"],
+                payload=result["payload"] if include_payload else None,
                 score=normalized_values[value_index],
             )
             for value_index, result in enumerate(result_values)
@@ -269,6 +280,7 @@ class LanceDBAdapter(VectorDBInterface):
         query_texts: List[str],
         limit: Optional[int] = None,
         with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         query_vectors = await self.embedding_engine.embed_text(query_texts)
 
@@ -279,6 +291,7 @@ class LanceDBAdapter(VectorDBInterface):
                     query_vector=query_vector,
                     limit=limit,
                     with_vector=with_vectors,
+                    include_payload=include_payload,
                 )
                 for query_vector in query_vectors
             ]

cognee/infrastructure/databases/vector/models/ScoredResult.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 from uuid import UUID
 from pydantic import BaseModel
 
@@ -12,10 +12,10 @@ class ScoredResult(BaseModel):
     - id (UUID): Unique identifier for the scored result.
     - score (float): The score associated with the result, where a lower score indicates a
     better outcome.
-    - payload (Dict[str, Any]): Additional information related to the score, stored as
+    - payload (Optional[Dict[str, Any]]): Additional information related to the score, stored as
     key-value pairs in a dictionary.
     """
 
     id: UUID
     score: float  # Lower score is better
-    payload: Dict[str, Any]
+    payload: Optional[Dict[str, Any]] = None

cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

@@ -301,6 +301,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         query_vector: Optional[List[float]] = None,
         limit: Optional[int] = 15,
         with_vector: bool = False,
+        include_payload: bool = False,
     ) -> List[ScoredResult]:
         if query_text is None and query_vector is None:
             raise MissingQueryParameterError()
@@ -324,10 +325,16 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         # NOTE: This needs to be initialized in case search doesn't return a value
         closest_items = []
 
+        # Note: Exclude payload from returned columns if not needed to optimize performance
+        select_columns = (
+            [PGVectorDataPoint]
+            if include_payload
+            else [PGVectorDataPoint.c.id, PGVectorDataPoint.c.vector]
+        )
         # Use async session to connect to the database
         async with self.get_async_session() as session:
             query = select(
-                PGVectorDataPoint,
+                *select_columns,
                 PGVectorDataPoint.c.vector.cosine_distance(query_vector).label("similarity"),
             ).order_by("similarity")
 
@@ -344,7 +351,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
             vector_list.append(
                 {
                     "id": parse_id(str(vector.id)),
-                    "payload": vector.payload,
+                    "payload": vector.payload if include_payload else None,
                     "_distance": vector.similarity,
                 }
             )
@@ -359,7 +366,11 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
 
         # Create and return ScoredResult objects
         return [
-            ScoredResult(id=row.get("id"), payload=row.get("payload"), score=row.get("score"))
+            ScoredResult(
+                id=row.get("id"),
+                payload=row.get("payload") if include_payload else None,
+                score=row.get("score"),
+            )
             for row in vector_list
         ]
 
@@ -369,6 +380,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         query_texts: List[str],
         limit: int = None,
         with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         query_vectors = await self.embedding_engine.embed_text(query_texts)
 
@@ -379,6 +391,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
                     query_vector=query_vector,
                     limit=limit,
                     with_vector=with_vectors,
+                    include_payload=include_payload,
                 )
                 for query_vector in query_vectors
             ]

cognee/infrastructure/databases/vector/pgvector/PGVectorDatasetDatabaseHandler.py

@@ -0,0 +1,86 @@
+from uuid import UUID
+from typing import Optional
+
+from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
+from cognee.infrastructure.databases.vector.pgvector.create_db_and_tables import delete_pg_database
+from cognee.modules.users.models import User
+from cognee.modules.users.models import DatasetDatabase
+from cognee.infrastructure.databases.vector import get_vectordb_config
+from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
+
+
+class PGVectorDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    """
+    Handler for interacting with PGVector Dataset databases.
+    """
+
+    @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 != "pgvector":
+            raise ValueError(
+                "PGVectorDatasetDatabaseHandler can only be used with PGVector vector database provider."
+            )
+
+        vector_db_name = f"{dataset_id}"
+
+        new_vector_config = {
+            "vector_database_provider": vector_config.vector_db_provider,
+            "vector_database_url": vector_config.vector_db_url,
+            "vector_database_name": vector_db_name,
+            "vector_database_connection_info": {
+                "port": vector_config.vector_db_port,
+                "host": vector_config.vector_db_host,
+            },
+            "vector_dataset_database_handler": "pgvector",
+        }
+
+        from .create_db_and_tables import create_pg_database
+
+        await create_pg_database(
+            {
+                "vector_db_provider": new_vector_config["vector_database_provider"],
+                "vector_db_url": new_vector_config["vector_database_url"],
+                "vector_db_name": new_vector_config["vector_database_name"],
+                "vector_db_port": new_vector_config["vector_database_connection_info"]["port"],
+                "vector_db_key": "",
+                "vector_db_username": vector_config.vector_db_username,
+                "vector_db_password": vector_config.vector_db_password,
+                "vector_db_host": new_vector_config["vector_database_connection_info"]["host"],
+                "vector_dataset_database_handler": "pgvector",
+            }
+        )
+
+        return new_vector_config
+
+    @classmethod
+    async def resolve_dataset_connection_info(
+        cls, dataset_database: DatasetDatabase
+    ) -> DatasetDatabase:
+        vector_config = get_vectordb_config()
+        # Note: For PGVector, we use the vector DB username/password from configuration so it's never stored in the DB
+        dataset_database.vector_database_connection_info["username"] = (
+            vector_config.vector_db_username
+        )
+        dataset_database.vector_database_connection_info["password"] = (
+            vector_config.vector_db_password
+        )
+        return dataset_database
+
+    @classmethod
+    async def delete_dataset(cls, dataset_database: DatasetDatabase):
+        dataset_database = await cls.resolve_dataset_connection_info(dataset_database)
+        vector_engine = create_vector_engine(
+            vector_db_provider=dataset_database.vector_database_provider,
+            vector_db_url=dataset_database.vector_database_url,
+            vector_db_name=dataset_database.vector_database_name,
+            vector_db_port=dataset_database.vector_database_connection_info["port"],
+            vector_db_username=dataset_database.vector_database_connection_info["username"],
+            vector_db_password=dataset_database.vector_database_connection_info["password"],
+        )
+        # Prune data
+        await vector_engine.prune()
+
+        # Drop entire database
+        await delete_pg_database(dataset_database)

cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py

@@ -1,5 +1,10 @@
-from sqlalchemy import text
+from sqlalchemy import text, URL
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from cognee.modules.users.models import DatasetDatabase
 from ..get_vector_engine import get_vector_engine, get_vectordb_context_config
+from ...vector import get_vectordb_config
+from cognee.context_global_variables import backend_access_control_enabled
 
 
 async def create_db_and_tables():
@@ -7,6 +12,80 @@ async def create_db_and_tables():
     vector_config = get_vectordb_context_config()
     vector_engine = get_vector_engine()
 
-    if vector_config["vector_db_provider"] == "pgvector":
+    if vector_config["vector_db_provider"] == "pgvector" and not backend_access_control_enabled():
         async with vector_engine.engine.begin() as connection:
             await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
+
+
+async def create_pg_database(vector_config):
+    """
+    Create the necessary Postgres database, and the PGVector extension on it.
+    This is defined separately because the creation needs the latest vector config,
+    which is not yet saved in the vector config context variable.
+    """
+
+    from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
+
+    # Create a maintenance engine, used when creating new postgres databases.
+    # Database named "postgres" should always exist. We need this since the SQLAlchemy
+    # engine cannot directly execute queries without first connecting to a database.
+    maintenance_db_name = "postgres"
+    maintenance_db_url = URL.create(
+        "postgresql+asyncpg",
+        username=vector_config["vector_db_username"],
+        password=vector_config["vector_db_password"],
+        host=vector_config["vector_db_host"],
+        port=int(vector_config["vector_db_port"]),
+        database=maintenance_db_name,
+    )
+    maintenance_engine = create_async_engine(maintenance_db_url)
+
+    # Connect to maintenance db in order to create new database
+    # Make sure to execute CREATE DATABASE outside of transaction block, and set AUTOCOMMIT isolation level
+    connection = await maintenance_engine.connect()
+    connection = await connection.execution_options(isolation_level="AUTOCOMMIT")
+    await connection.execute(text(f'CREATE DATABASE "{vector_config["vector_db_name"]}";'))
+
+    # Clean up resources
+    await connection.close()
+
+    vector_engine = create_vector_engine(**vector_config)
+    async with vector_engine.engine.begin() as connection:
+        await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
+
+
+async def delete_pg_database(dataset_database: DatasetDatabase):
+    """
+    Delete the Postgres database that was created for the PGVector extension,
+    in case of using PGVector with multi-user mode.
+    """
+
+    vector_config = get_vectordb_config()
+    # Create a maintenance engine, used when creating new postgres databases.
+    # Database named "postgres" should always exist. We need this since the SQLAlchemy
+    # engine cannot drop a database to which it is connected.
+    maintenance_db_name = "postgres"
+    maintenance_db_url = URL.create(
+        "postgresql+asyncpg",
+        username=vector_config.vector_db_username,
+        password=vector_config.vector_db_password,
+        host=vector_config.vector_db_host,
+        port=int(vector_config.vector_db_port),
+        database=maintenance_db_name,
+    )
+    maintenance_engine = create_async_engine(maintenance_db_url)
+
+    connection = await maintenance_engine.connect()
+    connection = await connection.execution_options(isolation_level="AUTOCOMMIT")
+    # We first have to kill all active sessions on the database, then delete it
+    await connection.execute(
+        text(
+            "SELECT pg_terminate_backend(pid) "
+            "FROM pg_stat_activity "
+            "WHERE datname = :db AND pid <> pg_backend_pid()"
+        ),
+        {"db": dataset_database.vector_database_name},
+    )
+    await connection.execute(text(f'DROP DATABASE "{dataset_database.vector_database_name}";'))
+
+    await connection.close()

cognee/infrastructure/databases/vector/vector_db_interface.py

@@ -87,6 +87,7 @@ class VectorDBInterface(Protocol):
         query_vector: Optional[List[float]],
         limit: Optional[int],
         with_vector: bool = False,
+        include_payload: bool = False,
     ):
         """
         Perform a search in the specified collection using either a text query or a vector
@@ -103,6 +104,9 @@ class VectorDBInterface(Protocol):
             - limit (Optional[int]): The maximum number of results to return from the search.
             - with_vector (bool): Whether to return the vector representations with search
               results. (default False)
+            - include_payload (bool): Whether to include the payload data with search. Search is faster when set to False.
+              Payload contains metadata about the data point, useful for searches that are only based on embedding distances
+              like the RAG_COMPLETION search type, but not needed when search also contains graph data.
         """
         raise NotImplementedError
 
@@ -113,6 +117,7 @@ class VectorDBInterface(Protocol):
         query_texts: List[str],
         limit: Optional[int],
         with_vectors: bool = False,
+        include_payload: bool = False,
     ):
         """
         Perform a batch search using multiple text queries against a collection.
@@ -125,6 +130,9 @@ class VectorDBInterface(Protocol):
             - limit (Optional[int]): The maximum number of results to return for each query.
             - with_vectors (bool): Whether to include vector representations with search
               results. (default False)
+            - include_payload (bool): Whether to include the payload data with search. Search is faster when set to False.
+              Payload contains metadata about the data point, useful for searches that are only based on embedding distances
+              like the RAG_COMPLETION search type, but not needed when search also contains graph data.
         """
         raise NotImplementedError