From c9332908868b6a920ee974a78b1060bb039e7e1d Mon Sep 17 00:00:00 2001
From: vasilije <vas.markovic@gmail.com>
Date: Wed, 27 Aug 2025 15:49:38 +0200
Subject: [PATCH] fixes to CLI

---
 cognee/cli/commands/cognify_command.py |  4 +-
 cognee/shared/logging_utils.py         | 93 +++++++++++++++++++++++++-
 2 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/cognee/cli/commands/cognify_command.py b/cognee/cli/commands/cognify_command.py
index bbe07f35a..894bfcd37 100644
--- a/cognee/cli/commands/cognify_command.py
+++ b/cognee/cli/commands/cognify_command.py
@@ -87,12 +87,12 @@ After successful cognify processing, use `cognee search` to query the knowledge
             async def run_cognify():
                 try:
                     # Import chunker classes here
-                    from cognee.modules.chunking import TextChunker
+                    from cognee.modules.chunking.TextChunker import TextChunker
 
                     chunker_class = TextChunker  # Default
                     if args.chunker == "LangchainChunker":
                         try:
-                            from cognee.modules.chunking import LangchainChunker
+                            from cognee.modules.chunking.LangchainChunker import LangchainChunker
 
                             chunker_class = LangchainChunker
                         except ImportError:
diff --git a/cognee/shared/logging_utils.py b/cognee/shared/logging_utils.py
index 29ac1761a..338626387 100644
--- a/cognee/shared/logging_utils.py
+++ b/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
+        
+        # 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
 
-        # Suppress LiteLLM ERROR logging using standard logging
-        logging.getLogger("litellm").setLevel(logging.CRITICAL)
+        # 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
@@ -243,6 +272,66 @@ def setup_logging(log_level=None, name=None):
 
     # 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:
+                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."""