cognee 0.2.3.dev0__py3-none-any.whl → 0.2.4__py3-none-any.whl

cognee/shared/CodeGraphEntities.py

@@ -36,6 +36,7 @@ class ClassDefinition(DataPoint):
 class CodeFile(DataPoint):
     name: str
     file_path: str
+    language: Optional[str] = None  # e.g., 'python', 'javascript', 'java', etc.
     source_code: Optional[str] = None
     part_of: Optional[Repository] = None
     depends_on: Optional[List["ImportStatement"]] = []

cognee/shared/exceptions/exceptions.py

@@ -1,8 +1,8 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeValidationError
 from fastapi import status
 
 
-class IngestionError(CogneeApiError):
+class IngestionError(CogneeValidationError):
     def __init__(
         self,
         message: str = "Failed to load data.",

cognee/shared/logging_utils.py

@@ -15,14 +15,43 @@ from typing import Protocol
 # Configure external library logging
 def configure_external_library_logging():
     """Configure logging for external libraries to reduce verbosity"""
+    # Set environment variables to suppress LiteLLM logging
+    os.environ.setdefault("LITELLM_LOG", "ERROR")
+    os.environ.setdefault("LITELLM_SET_VERBOSE", "False")
+
     # Configure LiteLLM logging to reduce verbosity
     try:
         import litellm
 
+        # Disable verbose logging
         litellm.set_verbose = False
 
-        # Suppress LiteLLM ERROR logging using standard logging
-        logging.getLogger("litellm").setLevel(logging.CRITICAL)
+        # Set additional LiteLLM configuration
+        if hasattr(litellm, "suppress_debug_info"):
+            litellm.suppress_debug_info = True
+        if hasattr(litellm, "turn_off_message"):
+            litellm.turn_off_message = True
+        if hasattr(litellm, "_turn_on_debug"):
+            litellm._turn_on_debug = False
+
+        # Comprehensive logger suppression
+        loggers_to_suppress = [
+            "litellm",
+            "litellm.litellm_core_utils.logging_worker",
+            "litellm.litellm_core_utils",
+            "litellm.proxy",
+            "litellm.router",
+            "openai._base_client",
+            "LiteLLM",  # Capital case variant
+            "LiteLLM.core",
+            "LiteLLM.logging_worker",
+            "litellm.logging_worker",
+        ]
+
+        for logger_name in loggers_to_suppress:
+            logging.getLogger(logger_name).setLevel(logging.CRITICAL)
+            logging.getLogger(logger_name).disabled = True
+
     except ImportError:
         # LiteLLM not available, skip configuration
         pass
@@ -173,29 +202,17 @@ def log_database_configuration(logger):
     from cognee.infrastructure.databases.graph.config import get_graph_config
 
     try:
-        # Log relational database configuration
-        relational_config = get_relational_config()
-        if relational_config.db_provider == "postgres":
-            logger.info(f"Postgres host: {relational_config.db_host}:{relational_config.db_port}")
-        elif relational_config.db_provider == "sqlite":
-            logger.info(f"SQLite path: {relational_config.db_path}")
-
-        # Log vector database configuration
-        vector_config = get_vectordb_config()
-        if vector_config.vector_db_provider == "lancedb":
-            logger.info(f"Vector database path: {vector_config.vector_db_url}")
-        else:
-            logger.info(f"Vector database URL: {vector_config.vector_db_url}")
-
-        # Log graph database configuration
-        graph_config = get_graph_config()
-        if graph_config.graph_database_provider == "kuzu":
-            logger.info(f"Graph database path: {graph_config.graph_file_path}")
-        else:
-            logger.info(f"Graph database URL: {graph_config.graph_database_url}")
+        # Get base database directory path
+        from cognee.base_config import get_base_config
+
+        base_config = get_base_config()
+        databases_path = os.path.join(base_config.system_root_directory, "databases")
+
+        # Log concise database info
+        logger.info(f"Database storage: {databases_path}")
 
     except Exception as e:
-        logger.warning(f"Could not retrieve database configuration: {str(e)}")
+        logger.debug(f"Could not retrieve database configuration: {str(e)}")
 
 
 def cleanup_old_logs(logs_dir, max_files):
@@ -216,13 +233,22 @@ def cleanup_old_logs(logs_dir, max_files):
 
         # Remove old files that exceed the maximum
         if len(log_files) > max_files:
+            deleted_count = 0
             for old_file in log_files[max_files:]:
                 try:
                     old_file.unlink()
-                    logger.info(f"Deleted old log file: {old_file}")
+                    deleted_count += 1
+                    # Only log individual files in non-CLI mode
+                    if os.getenv("COGNEE_CLI_MODE") != "true":
+                        logger.info(f"Deleted old log file: {old_file}")
                 except Exception as e:
+                    # Always log errors
                     logger.error(f"Failed to delete old log file {old_file}: {e}")
 
+            # In CLI mode, show compact summary
+            if os.getenv("COGNEE_CLI_MODE") == "true" and deleted_count > 0:
+                logger.info(f"Cleaned up {deleted_count} old log files")
+
         return True
     except Exception as e:
         logger.error(f"Error cleaning up log files: {e}")
@@ -241,11 +267,81 @@ def setup_logging(log_level=None, name=None):
     """
     global _is_structlog_configured
 
+    # Regular detailed logging for non-CLI usage
     log_level = log_level if log_level else log_levels[os.getenv("LOG_LEVEL", "INFO")]
 
     # Configure external library logging early to suppress verbose output
     configure_external_library_logging()
 
+    # Add custom filter to suppress LiteLLM worker cancellation errors
+    class LiteLLMCancellationFilter(logging.Filter):
+        """Filter to suppress LiteLLM worker cancellation messages"""
+
+        def filter(self, record):
+            # Check if this is a LiteLLM-related logger
+            if hasattr(record, "name") and "litellm" in record.name.lower():
+                return False
+
+            # Check message content for cancellation errors
+            if hasattr(record, "msg") and record.msg:
+                msg_str = str(record.msg).lower()
+                if any(
+                    keyword in msg_str
+                    for keyword in [
+                        "loggingworker cancelled",
+                        "logging_worker.py",
+                        "cancellederror",
+                        "litellm:error",
+                    ]
+                ):
+                    return False
+
+            # Check formatted message
+            try:
+                if hasattr(record, "getMessage"):
+                    formatted_msg = record.getMessage().lower()
+                    if any(
+                        keyword in formatted_msg
+                        for keyword in [
+                            "loggingworker cancelled",
+                            "logging_worker.py",
+                            "cancellederror",
+                            "litellm:error",
+                        ]
+                    ):
+                        return False
+            except Exception:
+                pass
+
+            return True
+
+    # Apply the filter to root logger and specific loggers
+    cancellation_filter = LiteLLMCancellationFilter()
+    logging.getLogger().addFilter(cancellation_filter)
+    logging.getLogger("litellm").addFilter(cancellation_filter)
+
+    # Add custom filter to suppress LiteLLM worker cancellation errors
+    class LiteLLMFilter(logging.Filter):
+        def filter(self, record):
+            # Suppress LiteLLM worker cancellation errors
+            if hasattr(record, "msg") and isinstance(record.msg, str):
+                msg_lower = record.msg.lower()
+                if any(
+                    phrase in msg_lower
+                    for phrase in [
+                        "loggingworker cancelled",
+                        "cancellederror",
+                        "logging_worker.py",
+                        "loggingerror",
+                    ]
+                ):
+                    return False
+            return True
+
+    # Apply filter to root logger
+    litellm_filter = LiteLLMFilter()
+    logging.getLogger().addFilter(litellm_filter)
+
     def exception_handler(logger, method_name, event_dict):
         """Custom processor to handle uncaught exceptions."""
         # Check if there's an exc_info that needs to be processed
@@ -298,11 +394,6 @@ def setup_logging(log_level=None, name=None):
         # Hand back to the original hook → prints traceback and exits
         sys.__excepthook__(exc_type, exc_value, traceback)
 
-        logger.info("Want to learn more? Visit the Cognee documentation: https://docs.cognee.ai")
-        logger.info(
-            "Need help? Reach out to us on our Discord server: https://discord.gg/NQPKmU5CCg"
-        )
-
     # Install exception handlers
     sys.excepthook = handle_exception
 
@@ -380,18 +471,38 @@ def setup_logging(log_level=None, name=None):
     # Mark logging as configured
     _is_structlog_configured = True
 
+    from cognee.infrastructure.databases.relational.config import get_relational_config
+    from cognee.infrastructure.databases.vector.config import get_vectordb_config
+    from cognee.infrastructure.databases.graph.config import get_graph_config
+
+    graph_config = get_graph_config()
+    vector_config = get_vectordb_config()
+    relational_config = get_relational_config()
+
+    try:
+        # Get base database directory path
+        from cognee.base_config import get_base_config
+
+        base_config = get_base_config()
+        databases_path = os.path.join(base_config.system_root_directory, "databases")
+    except Exception as e:
+        raise ValueError from e
+
     # Get a configured logger and log system information
     logger = structlog.get_logger(name if name else __name__)
+    # Detailed initialization for regular usage
     logger.info(
         "Logging initialized",
         python_version=PYTHON_VERSION,
         structlog_version=STRUCTLOG_VERSION,
         cognee_version=COGNEE_VERSION,
         os_info=OS_INFO,
+        database_path=databases_path,
+        graph_database_name=graph_config.graph_database_name,
+        vector_config=vector_config.vector_db_provider,
+        relational_config=relational_config.db_name,
     )
 
-    logger.info("Want to learn more? Visit the Cognee documentation: https://docs.cognee.ai")
-
     # Log database configuration
     log_database_configuration(logger)
 

cognee/shared/utils.py

@@ -3,7 +3,6 @@
 import os
 import requests
 from datetime import datetime, timezone
-import networkx as nx
 import matplotlib.pyplot as plt
 import http.server
 import socketserver

cognee/tasks/completion/exceptions/exceptions.py

@@ -1,11 +1,11 @@
-from cognee.exceptions import CogneeApiError
+from cognee.exceptions import CogneeValidationError
 from fastapi import status
 
 
-class NoRelevantDataError(CogneeApiError):
+class NoRelevantDataError(CogneeValidationError):
     """
     Represents an error when no relevant data is found during a search. This class is a
-    subclass of CogneeApiError.
+    subclass of CogneeValidationError.
 
     Public methods:
 

cognee/tasks/documents/classify_documents.py

@@ -10,6 +10,7 @@ from cognee.modules.data.processing.document_types import (
 )
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.engine.utils.generate_node_id import generate_node_id
+from cognee.tasks.documents.exceptions import WrongDataDocumentInputError
 
 EXTENSION_TO_DOCUMENT_CLASS = {
     "pdf": PdfDocument,  # Text documents
@@ -111,6 +112,9 @@ async def classify_documents(data_documents: list[Data]) -> list[Document]:
         - list[Document]: A list of Document objects created based on the classified data
           documents.
     """
+    if not isinstance(data_documents, list):
+        raise WrongDataDocumentInputError("data_documents")
+
     documents = []
     for data_item in data_documents:
         document = EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](

cognee/tasks/documents/exceptions/__init__.py

@@ -0,0 +1,11 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various data errors
+"""
+
+from .exceptions import (
+    WrongDataDocumentInputError,
+    InvalidChunkSizeError,
+    InvalidChunkerError,
+)

cognee/tasks/documents/exceptions/exceptions.py

@@ -0,0 +1,36 @@
+from cognee.exceptions import (
+    CogneeValidationError,
+    CogneeConfigurationError,
+)
+from fastapi import status
+
+
+class WrongDataDocumentInputError(CogneeValidationError):
+    """Raised when a wrong data document is provided."""
+
+    def __init__(
+        self,
+        field: str,
+        name: str = "WrongDataDocumentInputError",
+        status_code: int = status.HTTP_422_UNPROCESSABLE_ENTITY,
+    ):
+        message = f"Missing of invalid parameter: '{field}'."
+        super().__init__(message, name, status_code)
+
+
+class InvalidChunkSizeError(CogneeValidationError):
+    def __init__(self, value):
+        super().__init__(
+            message=f"max_chunk_size must be a positive integer (got {value}).",
+            name="InvalidChunkSizeError",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
+class InvalidChunkerError(CogneeValidationError):
+    def __init__(self):
+        super().__init__(
+            message="chunker must be a valid Chunker class.",
+            name="InvalidChunkerError",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )

cognee/tasks/documents/extract_chunks_from_documents.py

@@ -8,6 +8,7 @@ from cognee.modules.data.models import Data
 from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.chunking.TextChunker import TextChunker
 from cognee.modules.chunking.Chunker import Chunker
+from cognee.tasks.documents.exceptions import InvalidChunkSizeError, InvalidChunkerError
 
 
 async def update_document_token_count(document_id: UUID, token_count: int) -> None:
@@ -37,6 +38,13 @@ async def extract_chunks_from_documents(
         - The `read` method of the `Document` class must be implemented to support the chunking operation.
         - The `chunker` parameter determines the chunking logic and should align with the document type.
     """
+    if not isinstance(max_chunk_size, int) or max_chunk_size <= 0:
+        raise InvalidChunkSizeError(max_chunk_size)
+    if not isinstance(chunker, type):
+        raise InvalidChunkerError()
+    if not hasattr(chunker, "read"):
+        raise InvalidChunkerError()
+
     for document in documents:
         document_token_count = 0
 
@@ -48,5 +56,3 @@ async def extract_chunks_from_documents(
             yield document_chunk
 
         await update_document_token_count(document.id, document_token_count)
-
-        # todo rita

cognee/tasks/graph/exceptions/__init__.py

@@ -0,0 +1,12 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various data errors
+"""
+
+from .exceptions import (
+    InvalidDataChunksError,
+    InvalidGraphModelError,
+    InvalidOntologyAdapterError,
+    InvalidChunkGraphInputError,
+)

cognee/tasks/graph/exceptions/exceptions.py

@@ -0,0 +1,41 @@
+from cognee.exceptions import (
+    CogneeValidationError,
+    CogneeConfigurationError,
+)
+from fastapi import status
+
+
+class InvalidDataChunksError(CogneeValidationError):
+    def __init__(self, detail: str):
+        super().__init__(
+            message=f"Invalid data_chunks: {detail}",
+            name="InvalidDataChunksError",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
+class InvalidGraphModelError(CogneeValidationError):
+    def __init__(self, got):
+        super().__init__(
+            message=f"graph_model must be a subclass of BaseModel (got {got}).",
+            name="InvalidGraphModelError",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
+class InvalidOntologyAdapterError(CogneeConfigurationError):
+    def __init__(self, got):
+        super().__init__(
+            message=f"ontology_adapter lacks required interface (got {got}).",
+            name="InvalidOntologyAdapterError",
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+
+
+class InvalidChunkGraphInputError(CogneeValidationError):
+    def __init__(self, detail: str):
+        super().__init__(
+            message=f"Invalid chunk inputs or LLM Chunkgraphs: {detail}",
+            name="InvalidChunkGraphInputError",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )

cognee/tasks/graph/extract_graph_from_data.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Type, List
+from typing import Type, List, Optional
 from pydantic import BaseModel
 
 from cognee.infrastructure.databases.graph import get_graph_engine
@@ -12,6 +12,12 @@ from cognee.modules.graph.utils import (
 )
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.tasks.graph.exceptions import (
+    InvalidGraphModelError,
+    InvalidDataChunksError,
+    InvalidChunkGraphInputError,
+    InvalidOntologyAdapterError,
+)
 
 
 async def integrate_chunk_graphs(
@@ -21,6 +27,20 @@ async def integrate_chunk_graphs(
     ontology_adapter: OntologyResolver,
 ) -> List[DocumentChunk]:
     """Updates DocumentChunk objects, integrates data points and edges into databases."""
+
+    if not isinstance(data_chunks, list) or not isinstance(chunk_graphs, list):
+        raise InvalidChunkGraphInputError("data_chunks and chunk_graphs must be lists.")
+    if len(data_chunks) != len(chunk_graphs):
+        raise InvalidChunkGraphInputError(
+            f"length mismatch: {len(data_chunks)} chunks vs {len(chunk_graphs)} graphs."
+        )
+    if not isinstance(graph_model, type) or not issubclass(graph_model, BaseModel):
+        raise InvalidGraphModelError(graph_model)
+    if ontology_adapter is None or not hasattr(ontology_adapter, "get_subgraph"):
+        raise InvalidOntologyAdapterError(
+            type(ontology_adapter).__name__ if ontology_adapter else "None"
+        )
+
     graph_engine = await get_graph_engine()
 
     if graph_model is not KnowledgeGraph:
@@ -51,12 +71,24 @@ async def extract_graph_from_data(
     data_chunks: List[DocumentChunk],
     graph_model: Type[BaseModel],
     ontology_adapter: OntologyResolver = None,
+    custom_prompt: Optional[str] = None,
 ) -> List[DocumentChunk]:
     """
     Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model.
     """
+
+    if not isinstance(data_chunks, list) or not data_chunks:
+        raise InvalidDataChunksError("must be a non-empty list of DocumentChunk.")
+    if not all(hasattr(c, "text") for c in data_chunks):
+        raise InvalidDataChunksError("each chunk must have a 'text' attribute")
+    if not isinstance(graph_model, type) or not issubclass(graph_model, BaseModel):
+        raise InvalidGraphModelError(graph_model)
+
     chunk_graphs = await asyncio.gather(
-        *[LLMGateway.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
+        *[
+            LLMGateway.extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt)
+            for chunk in data_chunks
+        ]
     )
 
     # Note: Filter edges with missing source or target nodes

cognee/tasks/ingestion/exceptions/__init__.py

@@ -0,0 +1,8 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various application errors,
+such as System, Validation, Configuration or TransientErrors
+"""
+
+from .exceptions import S3FileSystemNotFoundError

cognee/tasks/ingestion/exceptions/exceptions.py

@@ -0,0 +1,12 @@
+from cognee.exceptions import CogneeSystemError
+from fastapi import status
+
+
+class S3FileSystemNotFoundError(CogneeSystemError):
+    def __init__(
+        self,
+        name: str = "S3FileSystemNotFoundError",
+        status_code: int = status.HTTP_500_INTERNAL_SERVER_ERROR,
+    ):
+        message = "Could not find S3FileSystem."
+        super().__init__(message, name, status_code)

cognee/tasks/ingestion/resolve_data_directories.py

@@ -1,6 +1,9 @@
 import os
 from urllib.parse import urlparse
 from typing import List, Union, BinaryIO
+
+from cognee.tasks.ingestion.exceptions import S3FileSystemNotFoundError
+from cognee.exceptions import CogneeSystemError
 from cognee.infrastructure.files.storage.s3_config import get_s3_config
 
 
@@ -54,6 +57,8 @@ async def resolve_data_directories(
                             else:
                                 s3_files.append(key)
                     resolved_data.extend(s3_files)
+                else:
+                    raise S3FileSystemNotFoundError()
 
             elif os.path.isdir(item):  # If it's a directory
                 if include_subdirectories:

cognee/tasks/repo_processor/get_local_dependencies.py

@@ -180,6 +180,7 @@ async def get_local_script_dependencies(
             name=file_path_relative_to_repo,
             source_code=source_code,
             file_path=script_path,
+            language="python",
         )
         return code_file_node
 
@@ -188,6 +189,7 @@ async def get_local_script_dependencies(
         name=file_path_relative_to_repo,
         source_code=None,
         file_path=script_path,
+        language="python",
     )
 
     async for part in extract_code_parts(source_code_tree.root_node, script_path=script_path):