cognee 0.5.0__py3-none-any.whl → 0.5.0.dev0__py3-none-any.whl

cognee/infrastructure/databases/graph/neo4j_driver/Neo4jAuraDevDatasetDatabaseHandler.py

@@ -1,168 +0,0 @@
-import os
-import asyncio
-import requests
-import base64
-import hashlib
-from uuid import UUID
-from typing import Optional
-from cryptography.fernet import Fernet
-
-from cognee.infrastructure.databases.graph import get_graph_config
-from cognee.modules.users.models import User, DatasetDatabase
-from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
-
-
-class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
-    """
-    Handler for a quick development PoC integration of Cognee multi-user and permission mode with Neo4j Aura databases.
-    This handler creates a new Neo4j Aura instance for each Cognee dataset created.
-
-    Improvements needed to be production ready:
-    - Secret management for client credentials, currently secrets are encrypted and stored in the Cognee relational database,
-      a secret manager or a similar system should be used instead.
-
-    Quality of life improvements:
-    - Allow configuration of different Neo4j Aura plans and regions.
-    - Requests should be made async, currently a blocking requests library is used.
-    """
-
-    @classmethod
-    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
-        """
-        Create a new Neo4j Aura instance for the dataset. Return connection info that will be mapped to the dataset.
-
-        Args:
-            dataset_id: Dataset UUID
-            user: User object who owns the dataset and is making the request
-
-        Returns:
-            dict: Connection details for the created Neo4j instance
-
-        """
-        graph_config = get_graph_config()
-
-        if graph_config.graph_database_provider != "neo4j":
-            raise ValueError(
-                "Neo4jAuraDevDatasetDatabaseHandler can only be used with Neo4j graph database provider."
-            )
-
-        graph_db_name = f"{dataset_id}"
-
-        # Client credentials and encryption
-        client_id = os.environ.get("NEO4J_CLIENT_ID", None)
-        client_secret = os.environ.get("NEO4J_CLIENT_SECRET", None)
-        tenant_id = os.environ.get("NEO4J_TENANT_ID", None)
-        encryption_env_key = os.environ.get("NEO4J_ENCRYPTION_KEY", "test_key")
-        encryption_key = base64.urlsafe_b64encode(
-            hashlib.sha256(encryption_env_key.encode()).digest()
-        )
-        cipher = Fernet(encryption_key)
-
-        if client_id is None or client_secret is None or tenant_id is None:
-            raise ValueError(
-                "NEO4J_CLIENT_ID, NEO4J_CLIENT_SECRET, and NEO4J_TENANT_ID environment variables must be set to use Neo4j Aura DatasetDatabase Handling."
-            )
-
-        # Make the request with HTTP Basic Auth
-        def get_aura_token(client_id: str, client_secret: str) -> dict:
-            url = "https://api.neo4j.io/oauth/token"
-            data = {"grant_type": "client_credentials"}  # sent as application/x-www-form-urlencoded
-
-            resp = requests.post(url, data=data, auth=(client_id, client_secret))
-            resp.raise_for_status()  # raises if the request failed
-            return resp.json()
-
-        resp = get_aura_token(client_id, client_secret)
-
-        url = "https://api.neo4j.io/v1/instances"
-
-        headers = {
-            "accept": "application/json",
-            "Authorization": f"Bearer {resp['access_token']}",
-            "Content-Type": "application/json",
-        }
-
-        # TODO: Maybe we can allow **kwargs parameter forwarding for cases like these
-        #       Too allow different configurations between datasets
-        payload = {
-            "version": "5",
-            "region": "europe-west1",
-            "memory": "1GB",
-            "name": graph_db_name[
-                0:29
-            ],  # TODO: Find better name to name Neo4j instance within 30 character limit
-            "type": "professional-db",
-            "tenant_id": tenant_id,
-            "cloud_provider": "gcp",
-        }
-
-        response = requests.post(url, headers=headers, json=payload)
-
-        graph_db_name = "neo4j"  # Has to be 'neo4j' for Aura
-        graph_db_url = response.json()["data"]["connection_url"]
-        graph_db_key = resp["access_token"]
-        graph_db_username = response.json()["data"]["username"]
-        graph_db_password = response.json()["data"]["password"]
-
-        async def _wait_for_neo4j_instance_provisioning(instance_id: str, headers: dict):
-            # Poll until the instance is running
-            status_url = f"https://api.neo4j.io/v1/instances/{instance_id}"
-            status = ""
-            for attempt in range(30):  # Try for up to ~5 minutes
-                status_resp = requests.get(
-                    status_url, headers=headers
-                )  # TODO: Use async requests with httpx
-                status = status_resp.json()["data"]["status"]
-                if status.lower() == "running":
-                    return
-                await asyncio.sleep(10)
-            raise TimeoutError(
-                f"Neo4j instance '{graph_db_name}' did not become ready within 5 minutes. Status: {status}"
-            )
-
-        instance_id = response.json()["data"]["id"]
-        await _wait_for_neo4j_instance_provisioning(instance_id, headers)
-
-        encrypted_db_password_bytes = cipher.encrypt(graph_db_password.encode())
-        encrypted_db_password_string = encrypted_db_password_bytes.decode()
-
-        return {
-            "graph_database_name": graph_db_name,
-            "graph_database_url": graph_db_url,
-            "graph_database_provider": "neo4j",
-            "graph_database_key": graph_db_key,
-            "graph_dataset_database_handler": "neo4j_aura_dev",
-            "graph_database_connection_info": {
-                "graph_database_username": graph_db_username,
-                "graph_database_password": encrypted_db_password_string,
-            },
-        }
-
-    @classmethod
-    async def resolve_dataset_connection_info(
-        cls, dataset_database: DatasetDatabase
-    ) -> DatasetDatabase:
-        """
-        Resolve and decrypt connection info for the Neo4j dataset database.
-        In this case, decrypt the password stored in the database.
-
-        Args:
-            dataset_database: DatasetDatabase instance containing encrypted connection info.
-        """
-        encryption_env_key = os.environ.get("NEO4J_ENCRYPTION_KEY", "test_key")
-        encryption_key = base64.urlsafe_b64encode(
-            hashlib.sha256(encryption_env_key.encode()).digest()
-        )
-        cipher = Fernet(encryption_key)
-        graph_db_password = cipher.decrypt(
-            dataset_database.graph_database_connection_info["graph_database_password"].encode()
-        ).decode()
-
-        dataset_database.graph_database_connection_info["graph_database_password"] = (
-            graph_db_password
-        )
-        return dataset_database
-
-    @classmethod
-    async def delete_dataset(cls, dataset_database: DatasetDatabase):
-        pass

cognee/infrastructure/databases/utils/get_graph_dataset_database_handler.py

@@ -1,10 +0,0 @@
-from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
-
-
-def get_graph_dataset_database_handler(dataset_database: DatasetDatabase) -> dict:
-    from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
-        supported_dataset_database_handlers,
-    )
-
-    handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler]
-    return handler

cognee/infrastructure/databases/utils/get_vector_dataset_database_handler.py

@@ -1,10 +0,0 @@
-from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
-
-
-def get_vector_dataset_database_handler(dataset_database: DatasetDatabase) -> dict:
-    from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
-        supported_dataset_database_handlers,
-    )
-
-    handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler]
-    return handler

cognee/infrastructure/databases/utils/resolve_dataset_database_connection_info.py

@@ -1,30 +0,0 @@
-from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
-    get_graph_dataset_database_handler,
-)
-from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
-    get_vector_dataset_database_handler,
-)
-from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
-
-
-async def resolve_dataset_database_connection_info(
-    dataset_database: DatasetDatabase,
-) -> DatasetDatabase:
-    """
-    Resolve the connection info for the given DatasetDatabase instance.
-    Resolve both vector and graph database connection info and return the updated DatasetDatabase instance.
-
-    Args:
-        dataset_database: DatasetDatabase instance
-    Returns:
-        DatasetDatabase instance with resolved connection info
-    """
-    vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
-    graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
-    dataset_database = await vector_dataset_database_handler[
-        "handler_instance"
-    ].resolve_dataset_connection_info(dataset_database)
-    dataset_database = await graph_dataset_database_handler[
-        "handler_instance"
-    ].resolve_dataset_connection_info(dataset_database)
-    return dataset_database

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

@@ -1,50 +0,0 @@
-import os
-from uuid import UUID
-from typing import Optional
-
-from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
-from cognee.modules.users.models import User
-from cognee.modules.users.models import DatasetDatabase
-from cognee.base_config import get_base_config
-from cognee.infrastructure.databases.vector import get_vectordb_config
-from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
-
-
-class LanceDBDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
-    """
-    Handler for interacting with LanceDB Dataset databases.
-    """
-
-    @classmethod
-    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
-        vector_config = get_vectordb_config()
-        base_config = get_base_config()
-
-        if vector_config.vector_db_provider != "lancedb":
-            raise ValueError(
-                "LanceDBDatasetDatabaseHandler can only be used with LanceDB vector database provider."
-            )
-
-        databases_directory_path = os.path.join(
-            base_config.system_root_directory, "databases", str(user.id)
-        )
-
-        vector_db_name = f"{dataset_id}.lance.db"
-
-        return {
-            "vector_database_provider": vector_config.vector_db_provider,
-            "vector_database_url": os.path.join(databases_directory_path, vector_db_name),
-            "vector_database_key": vector_config.vector_db_key,
-            "vector_database_name": vector_db_name,
-            "vector_dataset_database_handler": "lancedb",
-        }
-
-    @classmethod
-    async def delete_dataset(cls, dataset_database: DatasetDatabase):
-        vector_engine = create_vector_engine(
-            vector_db_provider=dataset_database.vector_database_provider,
-            vector_db_url=dataset_database.vector_database_url,
-            vector_db_key=dataset_database.vector_database_key,
-            vector_db_name=dataset_database.vector_database_name,
-        )
-        await vector_engine.prune()

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/__init__.py

@@ -1,5 +0,0 @@
-"""Bedrock LLM adapter module."""
-
-from .adapter import BedrockAdapter
-
-__all__ = ["BedrockAdapter"]

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/bedrock/adapter.py

@@ -1,153 +0,0 @@
-import litellm
-import instructor
-from typing import Type
-from pydantic import BaseModel
-from litellm.exceptions import ContentPolicyViolationError
-from instructor.exceptions import InstructorRetryException
-
-from cognee.infrastructure.llm.LLMGateway import LLMGateway
-from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
-    LLMInterface,
-)
-from cognee.infrastructure.llm.exceptions import (
-    ContentPolicyFilterError,
-    MissingSystemPromptPathError,
-)
-from cognee.infrastructure.files.storage.s3_config import get_s3_config
-from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
-    rate_limit_async,
-    rate_limit_sync,
-    sleep_and_retry_async,
-    sleep_and_retry_sync,
-)
-from cognee.modules.observability.get_observe import get_observe
-
-observe = get_observe()
-
-
-class BedrockAdapter(LLMInterface):
-    """
-    Adapter for AWS Bedrock API with support for three authentication methods:
-    1. API Key (Bearer Token)
-    2. AWS Credentials (access key + secret key)
-    3. AWS Profile (boto3 credential chain)
-    """
-
-    name = "Bedrock"
-    model: str
-    api_key: str
-    default_instructor_mode = "json_schema_mode"
-
-    MAX_RETRIES = 5
-
-    def __init__(
-        self,
-        model: str,
-        api_key: str = None,
-        max_completion_tokens: int = 16384,
-        streaming: bool = False,
-        instructor_mode: str = None,
-    ):
-        self.instructor_mode = instructor_mode if instructor_mode else self.default_instructor_mode
-
-        self.aclient = instructor.from_litellm(
-            litellm.acompletion, mode=instructor.Mode(self.instructor_mode)
-        )
-        self.client = instructor.from_litellm(litellm.completion)
-        self.model = model
-        self.api_key = api_key
-        self.max_completion_tokens = max_completion_tokens
-        self.streaming = streaming
-
-    def _create_bedrock_request(
-        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
-    ) -> dict:
-        """Create Bedrock request with authentication."""
-
-        request_params = {
-            "model": self.model,
-            "custom_llm_provider": "bedrock",
-            "drop_params": True,
-            "messages": [
-                {"role": "user", "content": text_input},
-                {"role": "system", "content": system_prompt},
-            ],
-            "response_model": response_model,
-            "max_retries": self.MAX_RETRIES,
-            "max_completion_tokens": self.max_completion_tokens,
-            "stream": self.streaming,
-        }
-
-        s3_config = get_s3_config()
-
-        # Add authentication parameters
-        if self.api_key:
-            request_params["api_key"] = self.api_key
-        elif s3_config.aws_access_key_id and s3_config.aws_secret_access_key:
-            request_params["aws_access_key_id"] = s3_config.aws_access_key_id
-            request_params["aws_secret_access_key"] = s3_config.aws_secret_access_key
-            if s3_config.aws_session_token:
-                request_params["aws_session_token"] = s3_config.aws_session_token
-        elif s3_config.aws_profile_name:
-            request_params["aws_profile_name"] = s3_config.aws_profile_name
-
-        if s3_config.aws_region:
-            request_params["aws_region_name"] = s3_config.aws_region
-
-        # Add optional parameters
-        if s3_config.aws_bedrock_runtime_endpoint:
-            request_params["aws_bedrock_runtime_endpoint"] = s3_config.aws_bedrock_runtime_endpoint
-
-        return request_params
-
-    @observe(as_type="generation")
-    @sleep_and_retry_async()
-    @rate_limit_async
-    async def acreate_structured_output(
-        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
-    ) -> BaseModel:
-        """Generate structured output from AWS Bedrock API."""
-
-        try:
-            request_params = self._create_bedrock_request(text_input, system_prompt, response_model)
-            return await self.aclient.chat.completions.create(**request_params)
-
-        except (
-            ContentPolicyViolationError,
-            InstructorRetryException,
-        ) as error:
-            if (
-                isinstance(error, InstructorRetryException)
-                and "content management policy" not in str(error).lower()
-            ):
-                raise error
-
-            raise ContentPolicyFilterError(
-                f"The provided input contains content that is not aligned with our content policy: {text_input}"
-            )
-
-    @observe
-    @sleep_and_retry_sync()
-    @rate_limit_sync
-    def create_structured_output(
-        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
-    ) -> BaseModel:
-        """Generate structured output from AWS Bedrock API (synchronous)."""
-
-        request_params = self._create_bedrock_request(text_input, system_prompt, response_model)
-        return self.client.chat.completions.create(**request_params)
-
-    def show_prompt(self, text_input: str, system_prompt: str) -> str:
-        """Format and display the prompt for a user query."""
-        if not text_input:
-            text_input = "No user input provided."
-        if not system_prompt:
-            raise MissingSystemPromptPathError()
-        system_prompt = LLMGateway.read_query_prompt(system_prompt)
-
-        formatted_prompt = (
-            f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
-            if system_prompt
-            else None
-        )
-        return formatted_prompt

cognee/memify_pipelines/create_triplet_embeddings.py

@@ -1,53 +0,0 @@
-from typing import Any
-
-from cognee import memify
-from cognee.context_global_variables import (
-    set_database_global_context_variables,
-)
-from cognee.exceptions import CogneeValidationError
-from cognee.modules.data.methods import get_authorized_existing_datasets
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.pipelines.tasks.task import Task
-from cognee.modules.users.models import User
-from cognee.tasks.memify.get_triplet_datapoints import get_triplet_datapoints
-from cognee.tasks.storage import index_data_points
-
-logger = get_logger("create_triplet_embeddings")
-
-
-async def create_triplet_embeddings(
-    user: User,
-    dataset: str = "main_dataset",
-    run_in_background: bool = False,
-    triplets_batch_size: int = 100,
-) -> dict[str, Any]:
-    dataset_to_write = await get_authorized_existing_datasets(
-        user=user, datasets=[dataset], permission_type="write"
-    )
-
-    if not dataset_to_write:
-        raise CogneeValidationError(
-            message=f"User does not have write access to dataset: {dataset}",
-            log=False,
-        )
-
-    await set_database_global_context_variables(
-        dataset_to_write[0].id, dataset_to_write[0].owner_id
-    )
-
-    extraction_tasks = [Task(get_triplet_datapoints, triplets_batch_size=triplets_batch_size)]
-
-    enrichment_tasks = [
-        Task(index_data_points, task_config={"batch_size": triplets_batch_size}),
-    ]
-
-    result = await memify(
-        extraction_tasks=extraction_tasks,
-        enrichment_tasks=enrichment_tasks,
-        dataset=dataset_to_write[0].id,
-        data=[{}],
-        user=user,
-        run_in_background=run_in_background,
-    )
-
-    return result

cognee/modules/engine/models/Triplet.py

@@ -1,9 +0,0 @@
-from cognee.infrastructure.engine import DataPoint
-
-
-class Triplet(DataPoint):
-    text: str
-    from_node_id: str
-    to_node_id: str
-
-    metadata: dict = {"index_fields": ["text"]}

cognee/modules/retrieval/register_retriever.py

@@ -1,10 +0,0 @@
-from typing import Type
-
-from .base_retriever import BaseRetriever
-from .registered_community_retrievers import registered_community_retrievers
-from ..search.types import SearchType
-
-
-def use_retriever(search_type: SearchType, retriever: Type[BaseRetriever]):
-    """Register a retriever class for a given search type."""
-    registered_community_retrievers[search_type] = retriever

cognee/modules/retrieval/registered_community_retrievers.py

@@ -1 +0,0 @@
-registered_community_retrievers = {}