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

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/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
@@ -71,6 +71,7 @@ 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.
@@ -84,7 +85,10 @@ async def extract_graph_from_data(
         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/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):

cognee/tasks/repo_processor/get_repo_file_dependencies.py

@@ -10,50 +10,80 @@ from cognee.infrastructure.engine import DataPoint
 from cognee.shared.CodeGraphEntities import CodeFile, Repository
 
 
-async def get_source_code_files(repo_path):
+async def get_source_code_files(repo_path, language_config: dict[str, list[str]] | None = None):
     """
-    Retrieve Python source code files from the specified repository path.
-
-    This function scans the given repository path for files that have the .py extension
-    while excluding test files and files within a virtual environment. It returns a list of
-    absolute paths to the source code files that are not empty.
+    Retrieve source code files from the specified repository path for multiple languages.
 
     Parameters:
     -----------
-
-        - repo_path: The file path to the repository to search for Python source files.
+        - repo_path: The file path to the repository to search for source files.
+        - language_config: dict mapping language names to file extensions, e.g.,
+            {'python': ['.py'], 'javascript': ['.js', '.jsx'], ...}
 
     Returns:
     --------
-
-        A list of absolute paths to .py files that contain source code, excluding empty
-        files, test files, and files from a virtual environment.
+        A list of (absolute_path, language) tuples for source code files.
     """
-    if not os.path.exists(repo_path):
-        return {}
-
-    py_files_paths = (
-        os.path.join(root, file)
-        for root, _, files in os.walk(repo_path)
-        for file in files
-        if (
-            file.endswith(".py")
-            and not file.startswith("test_")
-            and not file.endswith("_test")
-            and ".venv" not in file
-        )
-    )
 
-    source_code_files = set()
-    for file_path in py_files_paths:
-        file_path = os.path.abspath(file_path)
-
-        if os.path.getsize(file_path) == 0:
-            continue
+    def _get_language_from_extension(file, language_config):
+        for lang, exts in language_config.items():
+            for ext in exts:
+                if file.endswith(ext):
+                    return lang
+        return None
+
+    # Default config if not provided
+    if language_config is None:
+        language_config = {
+            "python": [".py"],
+            "javascript": [".js", ".jsx"],
+            "typescript": [".ts", ".tsx"],
+            "java": [".java"],
+            "csharp": [".cs"],
+            "go": [".go"],
+            "rust": [".rs"],
+            "cpp": [".cpp", ".c", ".h", ".hpp"],
+        }
 
-        source_code_files.add(file_path)
+    if not os.path.exists(repo_path):
+        return []
 
-    return list(source_code_files)
+    source_code_files = set()
+    for root, _, files in os.walk(repo_path):
+        for file in files:
+            lang = _get_language_from_extension(file, language_config)
+            if lang is None:
+                continue
+            # Exclude tests and common build/venv directories
+            excluded_dirs = {
+                ".venv",
+                "venv",
+                "env",
+                ".env",
+                "site-packages",
+                "node_modules",
+                "dist",
+                "build",
+                ".git",
+                "tests",
+                "test",
+            }
+            root_parts = set(os.path.normpath(root).split(os.sep))
+            base_name, _ext = os.path.splitext(file)
+            if (
+                base_name.startswith("test_")
+                or base_name.endswith("_test")  # catches Go's *_test.go and similar
+                or ".test." in file
+                or ".spec." in file
+                or (excluded_dirs & root_parts)
+            ):
+                continue
+            file_path = os.path.abspath(os.path.join(root, file))
+            if os.path.getsize(file_path) == 0:
+                continue
+            source_code_files.add((file_path, lang))
+
+    return sorted(list(source_code_files))
 
 
 def run_coroutine(coroutine_func, *args, **kwargs):
@@ -85,22 +115,23 @@ def run_coroutine(coroutine_func, *args, **kwargs):
 
 
 async def get_repo_file_dependencies(
-    repo_path: str, detailed_extraction: bool = False
+    repo_path: str, detailed_extraction: bool = False, supported_languages: list = None
 ) -> AsyncGenerator[DataPoint, None]:
     """
-    Generate a dependency graph for Python files in the given repository path.
+    Generate a dependency graph for source files (multi-language) in the given repository path.
 
     Check the validity of the repository path and yield a repository object followed by the
-    dependencies of Python files within that repository. Raise a FileNotFoundError if the
+    dependencies of source files within that repository. Raise a FileNotFoundError if the
     provided path does not exist. The extraction of detailed dependencies can be controlled
-    via the `detailed_extraction` argument.
+    via the `detailed_extraction` argument. Languages considered can be restricted via
+    the `supported_languages` argument.
 
     Parameters:
     -----------
 
-        - repo_path (str): The file path to the repository where Python files are located.
-        - detailed_extraction (bool): A flag indicating whether to perform a detailed
-          extraction of dependencies (default is False). (default False)
+        - repo_path (str): The file path to the repository to process.
+        - detailed_extraction (bool): Whether to perform a detailed extraction of code parts.
+        - supported_languages (list | None): Subset of languages to include; if None, use defaults.
     """
 
     if isinstance(repo_path, list) and len(repo_path) == 1:
@@ -109,7 +140,25 @@ async def get_repo_file_dependencies(
     if not os.path.exists(repo_path):
         raise FileNotFoundError(f"Repository path {repo_path} does not exist.")
 
-    source_code_files = await get_source_code_files(repo_path)
+    # Build language config from supported_languages
+    default_language_config = {
+        "python": [".py"],
+        "javascript": [".js", ".jsx"],
+        "typescript": [".ts", ".tsx"],
+        "java": [".java"],
+        "csharp": [".cs"],
+        "go": [".go"],
+        "rust": [".rs"],
+        "cpp": [".cpp", ".c", ".h", ".hpp"],
+    }
+    if supported_languages is not None:
+        language_config = {
+            k: v for k, v in default_language_config.items() if k in supported_languages
+        }
+    else:
+        language_config = default_language_config
+
+    source_code_files = await get_source_code_files(repo_path, language_config=language_config)
 
     repo = Repository(
         id=uuid5(NAMESPACE_OID, repo_path),
@@ -128,19 +177,42 @@ async def get_repo_file_dependencies(
         for chunk_number in range(number_of_chunks)
     ]
 
-    # Codegraph dependencies are not installed by default, so we import where we use them.
+    # Import dependency extractors for each language (Python for now, extend later)
     from cognee.tasks.repo_processor.get_local_dependencies import get_local_script_dependencies
+    import aiofiles
+    # TODO: Add other language extractors here
 
     for start_range, end_range in chunk_ranges:
-        # with ProcessPoolExecutor(max_workers=12) as executor:
-        tasks = [
-            get_local_script_dependencies(repo_path, file_path, detailed_extraction)
-            for file_path in source_code_files[start_range : end_range + 1]
-        ]
+        tasks = []
+        for file_path, lang in source_code_files[start_range : end_range + 1]:
+            # For now, only Python is supported; extend with other languages
+            if lang == "python":
+                tasks.append(
+                    get_local_script_dependencies(repo_path, file_path, detailed_extraction)
+                )
+            else:
+                # Placeholder: create a minimal CodeFile for other languages
+                async def make_codefile_stub(file_path=file_path, lang=lang):
+                    async with aiofiles.open(
+                        file_path, "r", encoding="utf-8", errors="replace"
+                    ) as f:
+                        source = await f.read()
+                    return CodeFile(
+                        id=uuid5(NAMESPACE_OID, file_path),
+                        name=os.path.relpath(file_path, repo_path),
+                        file_path=file_path,
+                        language=lang,
+                        source_code=source,
+                    )
+
+                tasks.append(make_codefile_stub())
 
         results: list[CodeFile] = await asyncio.gather(*tasks)
 
         for source_code_file in results:
             source_code_file.part_of = repo
-
+            if getattr(
+                source_code_file, "language", None
+            ) is None and source_code_file.file_path.endswith(".py"):
+                source_code_file.language = "python"
             yield source_code_file

cognee/tasks/storage/add_data_points.py

@@ -10,7 +10,37 @@ from cognee.tasks.storage.exceptions import (
 )
 
 
-async def add_data_points(data_points: List[DataPoint]) -> List[DataPoint]:
+async def add_data_points(
+    data_points: List[DataPoint], update_edge_collection: bool = True
+) -> List[DataPoint]:
+    """
+    Add a batch of data points to the graph database by extracting nodes and edges,
+    deduplicating them, and indexing them for retrieval.
+
+    This function parallelizes the graph extraction for each data point,
+    merges the resulting nodes and edges, and ensures uniqueness before
+    committing them to the underlying graph engine. It also updates the
+    associated retrieval indices for nodes and (optionally) edges.
+
+    Args:
+        data_points (List[DataPoint]):
+            A list of data points to process and insert into the graph.
+        update_edge_collection (bool, optional):
+            Whether to update the edge index after adding edges.
+            Defaults to True.
+
+    Returns:
+        List[DataPoint]:
+            The original list of data points after processing and insertion.
+
+    Side Effects:
+        - Calls `get_graph_from_model` concurrently for each data point.
+        - Deduplicates nodes and edges across all results.
+        - Updates the node index via `index_data_points`.
+        - Inserts nodes and edges into the graph engine.
+        - Optionally updates the edge index via `index_graph_edges`.
+    """
+
     if not isinstance(data_points, list):
         raise InvalidDataPointsInAddDataPointsError("data_points must be a list.")
     if not all(isinstance(dp, DataPoint) for dp in data_points):
@@ -48,7 +78,7 @@ async def add_data_points(data_points: List[DataPoint]) -> List[DataPoint]:
     await graph_engine.add_nodes(nodes)
     await graph_engine.add_edges(edges)
 
-    # This step has to happen after adding nodes and edges because we query the graph.
-    await index_graph_edges()
+    if update_edge_collection:
+        await index_graph_edges()
 
     return data_points

cognee/tests/integration/cli/__init__.py

@@ -0,0 +1,3 @@
+"""
+CLI integration tests package.
+"""