praisonaiagents 0.0.142__py3-none-any.whl → 0.0.143__py3-none-any.whl

praisonaiagents/telemetry/telemetry.py

@@ -9,6 +9,7 @@ import os
 import time
 import platform
 import hashlib
+import threading
 from typing import Dict, Any, Optional
 from datetime import datetime
 import logging
@@ -40,6 +41,15 @@ class MinimalTelemetry:
     - Can be disabled via environment variables
     """
     
+    # Common error phrases that indicate interpreter shutdown
+    _SHUTDOWN_ERROR_PHRASES = [
+        'cannot schedule new futures',
+        'interpreter shutdown',
+        'atexit after shutdown',
+        'event loop closed',
+        'runtime is shutting down'
+    ]
+    
     def __init__(self, enabled: bool = None):
         """
         Initialize the minimal telemetry collector.
@@ -55,6 +65,10 @@ class MinimalTelemetry:
             
         self.logger = logging.getLogger(__name__)
         
+        # Add shutdown tracking to prevent double shutdown
+        self._shutdown_complete = False
+        self._shutdown_lock = threading.Lock()
+        
         if not self.enabled:
             self.logger.debug("Telemetry is disabled")
             return
@@ -70,6 +84,8 @@ class MinimalTelemetry:
             "tool_calls": 0,
             "errors": 0,
         }
+        self._metrics_lock = threading.Lock()
+        self._max_timing_entries = 1000  # Limit to prevent memory leaks
         
         # Collect basic environment info (anonymous)
         self._environment = {
@@ -100,7 +116,7 @@ class MinimalTelemetry:
         try:
             from .. import __version__
             return __version__
-        except ImportError:
+        except (ImportError, KeyError, AttributeError):
             return "unknown"
     
     def track_agent_execution(self, agent_name: str = None, success: bool = True):
@@ -114,7 +130,8 @@ class MinimalTelemetry:
         if not self.enabled:
             return
             
-        self._metrics["agent_executions"] += 1
+        with self._metrics_lock:
+            self._metrics["agent_executions"] += 1
         
         # Send event to PostHog
         if self._posthog:
@@ -140,7 +157,8 @@ class MinimalTelemetry:
         if not self.enabled:
             return
             
-        self._metrics["task_completions"] += 1
+        with self._metrics_lock:
+            self._metrics["task_completions"] += 1
         
         # Send event to PostHog
         if self._posthog:
@@ -155,33 +173,60 @@ class MinimalTelemetry:
         
         self.logger.debug(f"Task completion tracked: success={success}")
     
-    def track_tool_usage(self, tool_name: str, success: bool = True):
+    def track_tool_usage(self, tool_name: str, success: bool = True, execution_time: float = None):
         """
-        Track tool usage event.
+        Track tool usage event with optional timing.
         
         Args:
             tool_name: Name of the tool being used
             success: Whether the tool call was successful
+            execution_time: Time in seconds the tool took to execute (optional)
         """
         if not self.enabled:
             return
             
-        self._metrics["tool_calls"] += 1
+        with self._metrics_lock:
+            self._metrics["tool_calls"] += 1
+            
+            # Add timing metrics if provided (with memory management)
+            if execution_time is not None:
+                if "tool_execution_times" not in self._metrics:
+                    self._metrics["tool_execution_times"] = []
+                
+                timing_list = self._metrics["tool_execution_times"]
+                timing_list.append({
+                    "tool_name": tool_name,
+                    "execution_time": execution_time,
+                    "success": success
+                })
+                
+                # Prevent memory leaks by limiting stored entries
+                if len(timing_list) > self._max_timing_entries:
+                    timing_list[:] = timing_list[-self._max_timing_entries:]
         
         # Send event to PostHog
         if self._posthog:
+            properties = {
+                'tool_name': tool_name,
+                'success': success,
+                'session_id': self.session_id
+            }
+            
+            # Include execution time if available
+            if execution_time is not None:
+                properties['execution_time'] = execution_time
+            
             self._posthog.capture(
                 distinct_id=self.session_id,
                 event='tool_usage',
-                properties={
-                    'tool_name': tool_name,
-                    'success': success,
-                    'session_id': self.session_id
-                }
+                properties=properties
             )
         
         # Only track tool name, not arguments or results
-        self.logger.debug(f"Tool usage tracked: {tool_name}, success={success}")
+        debug_msg = f"Tool usage tracked: {tool_name}, success={success}"
+        if execution_time is not None:
+            debug_msg += f", execution_time={execution_time:.3f}s"
+        self.logger.debug(debug_msg)
     
     def track_error(self, error_type: str = None):
         """
@@ -193,7 +238,8 @@ class MinimalTelemetry:
         if not self.enabled:
             return
             
-        self._metrics["errors"] += 1
+        with self._metrics_lock:
+            self._metrics["errors"] += 1
         
         # Send event to PostHog
         if self._posthog:
@@ -243,10 +289,13 @@ class MinimalTelemetry:
         if not self.enabled:
             return {"enabled": False}
             
+        with self._metrics_lock:
+            metrics_copy = self._metrics.copy()
+            
         return {
             "enabled": True,
             "session_id": self.session_id,
-            "metrics": self._metrics.copy(),
+            "metrics": metrics_copy,
             "environment": self._environment.copy(),
         }
     
@@ -281,28 +330,174 @@ class MinimalTelemetry:
                 pass
         
         # Reset counters
-        for key in self._metrics:
-            if isinstance(self._metrics[key], int):
-                self._metrics[key] = 0
+        with self._metrics_lock:
+            for key in self._metrics:
+                if isinstance(self._metrics[key], int):
+                    self._metrics[key] = 0
     
     def shutdown(self):
         """
         Shutdown telemetry and ensure all events are sent.
+        Forces proper cleanup of background threads to prevent hanging.
         """
         if not self.enabled:
             return
+        
+        # Use lock to prevent concurrent shutdown calls
+        with self._shutdown_lock:
+            if self._shutdown_complete:
+                return
+            self._shutdown_complete = True
             
         # Final flush
         self.flush()
         
         # Shutdown PostHog if available
-        if hasattr(self, '_posthog') and self._posthog:
+        posthog_client = getattr(self, '_posthog', None)
+        if posthog_client:
             try:
-                # Force a synchronous flush before shutdown
-                self._posthog.flush()
-                self._posthog.shutdown()
-            except:
-                pass
+                # Check if Python interpreter is shutting down
+                if self._is_interpreter_shutting_down():
+                    self.logger.debug("Interpreter shutting down, skipping PostHog operations")
+                    return
+                
+                # Use a timeout-based flush to prevent hanging
+                import threading
+                import time
+                import concurrent.futures
+                
+                # Use ThreadPoolExecutor for better control
+                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+                    flush_future = executor.submit(self._safe_flush_posthog, posthog_client)
+                    
+                    try:
+                        flush_future.result(timeout=5.0)  # 5 second timeout
+                        self.logger.debug("PostHog flush completed successfully")
+                    except concurrent.futures.TimeoutError:
+                        self.logger.warning("PostHog flush timed out")
+                        flush_future.cancel()
+                    except Exception as e:
+                        self.logger.error(f"PostHog flush failed: {e}")
+                
+                # Cleanup PostHog threads safely
+                self._shutdown_posthog_threads(posthog_client)
+                
+                # Standard shutdown - with interpreter shutdown check
+                if not self._is_interpreter_shutting_down():
+                    posthog_client.shutdown()
+                else:
+                    self.logger.debug("Skipping PostHog shutdown call due to interpreter shutdown")
+                
+            except Exception as e:
+                # Handle specific shutdown-related errors gracefully
+                if self._is_shutdown_related_error(e):
+                    self.logger.debug(f"PostHog shutdown prevented due to interpreter shutdown: {e}")
+                else:
+                    self.logger.error(f"Error during PostHog shutdown: {e}")
+            finally:
+                self._posthog = None
+    
+    def _is_shutdown_related_error(self, error: Exception) -> bool:
+        """
+        Check if an error is related to interpreter shutdown.
+        
+        Args:
+            error: The exception to check
+            
+        Returns:
+            True if the error is shutdown-related, False otherwise
+        """
+        error_msg = str(error).lower()
+        return any(phrase in error_msg for phrase in self._SHUTDOWN_ERROR_PHRASES)
+    
+    def _is_interpreter_shutting_down(self) -> bool:
+        """
+        Check if the Python interpreter is shutting down.
+        
+        Returns:
+            True if interpreter is shutting down, False otherwise
+        """
+        try:
+            import sys
+            
+            # Check if the interpreter is in shutdown mode
+            if hasattr(sys, 'is_finalizing') and sys.is_finalizing():
+                return True
+            
+            # Check if we can create new threads (fails during shutdown)
+            try:
+                test_thread = threading.Thread(target=lambda: None)
+                test_thread.daemon = True
+                test_thread.start()
+                test_thread.join(timeout=0.001)
+                return False
+            except (RuntimeError, threading.ThreadError):
+                return True
+                
+        except Exception:
+            # If we can't determine state, assume we're shutting down to be safe
+            return True
+    
+    def _safe_flush_posthog(self, posthog_client):
+        """Safely flush PostHog data with error handling."""
+        try:
+            # Skip flush if interpreter is shutting down
+            if self._is_interpreter_shutting_down():
+                self.logger.debug("Skipping PostHog flush due to interpreter shutdown")
+                return False
+                
+            posthog_client.flush()
+            return True
+        except Exception as e:
+            if self._is_shutdown_related_error(e):
+                self.logger.debug(f"PostHog flush prevented due to interpreter shutdown: {e}")
+            else:
+                self.logger.debug(f"PostHog flush error: {e}")
+            return False
+    
+    def _shutdown_posthog_threads(self, posthog_client):
+        """Safely shutdown PostHog background threads."""
+        try:
+            # Skip thread cleanup if interpreter is shutting down
+            if self._is_interpreter_shutting_down():
+                self.logger.debug("Skipping PostHog thread cleanup due to interpreter shutdown")
+                return
+                
+            # Access thread pool safely (fix double shutdown issue)
+            thread_pool = getattr(posthog_client, '_thread_pool', None)
+            if thread_pool:
+                try:
+                    # Single shutdown call with timeout
+                    if hasattr(thread_pool, 'shutdown'):
+                        thread_pool.shutdown(wait=False)
+                        # Wait briefly for graceful shutdown
+                        import time
+                        time.sleep(0.5)
+                except Exception as e:
+                    if self._is_shutdown_related_error(e):
+                        self.logger.debug(f"Thread pool shutdown prevented due to interpreter shutdown: {e}")
+                    else:
+                        self.logger.debug(f"Thread pool shutdown error: {e}")
+            
+            # Clean up consumer
+            consumer = getattr(posthog_client, '_consumer', None)
+            if consumer:
+                try:
+                    if hasattr(consumer, 'flush'):
+                        consumer.flush()
+                    if hasattr(consumer, 'shutdown'):
+                        consumer.shutdown()
+                except Exception as e:
+                    if self._is_shutdown_related_error(e):
+                        self.logger.debug(f"Consumer shutdown prevented due to interpreter shutdown: {e}")
+                    else:
+                        self.logger.debug(f"Consumer shutdown error: {e}")
+                    
+        except Exception as e:
+            if self._is_shutdown_related_error(e):
+                self.logger.debug(f"PostHog thread cleanup prevented due to interpreter shutdown: {e}")
+            else:
+                self.logger.debug(f"Error during PostHog thread cleanup: {e}")
 
 
 # Global telemetry instance
@@ -331,6 +526,41 @@ def disable_telemetry():
         _telemetry_instance = MinimalTelemetry(enabled=False)
 
 
+def force_shutdown_telemetry():
+    """
+    Force shutdown of telemetry system with comprehensive cleanup.
+    This function ensures proper termination of all background threads.
+    """
+    global _telemetry_instance
+    if _telemetry_instance:
+        _telemetry_instance.shutdown()
+        
+        # Additional cleanup - wait for all threads to finish
+        import threading
+        import time
+        
+        # Wait up to 3 seconds for any remaining threads to finish
+        max_wait = 3.0
+        start_time = time.time()
+        
+        while time.time() - start_time < max_wait:
+            # Check for any analytics/telemetry related threads
+            analytics_threads = [
+                t for t in threading.enumerate() 
+                if t != threading.current_thread() 
+                and not t.daemon
+                and any(keyword in t.name.lower() for keyword in ['posthog', 'analytics', 'telemetry', 'consumer'])
+            ]
+            
+            if not analytics_threads:
+                break
+                
+            time.sleep(0.1)
+        
+        # Reset the global instance
+        _telemetry_instance = None
+
+
 def enable_telemetry():
     """Programmatically enable telemetry (if not disabled by environment)."""
     global _telemetry_instance

praisonaiagents/tools/__init__.py

@@ -49,6 +49,20 @@ TOOL_MAPPINGS = {
     'analyze_data': ('.duckdb_tools', None),
     'duckdb_tools': ('.duckdb_tools', None),
     
+    # MongoDB Tools
+    'insert_document': ('.mongodb_tools', None),
+    'insert_documents': ('.mongodb_tools', None),
+    'find_documents': ('.mongodb_tools', None),
+    'update_document': ('.mongodb_tools', None),
+    'delete_document': ('.mongodb_tools', None),
+    'create_vector_index': ('.mongodb_tools', None),
+    'vector_search': ('.mongodb_tools', None),
+    'store_with_embedding': ('.mongodb_tools', None),
+    'text_search': ('.mongodb_tools', None),
+    'get_stats': ('.mongodb_tools', None),
+    'connect_mongodb': ('.mongodb_tools', None),
+    'mongodb_tools': ('.mongodb_tools', None),
+    
     # Shell Tools
     'execute_command': ('.shell_tools', None),
     'list_processes': ('.shell_tools', None),
@@ -180,12 +194,14 @@ def __getattr__(name: str) -> Any:
             'get_article', 'get_news_sources', 'get_articles_from_source', 'get_trending_topics',
             'scrape_page', 'extract_links', 'crawl', 'extract_text',
             'query', 'create_table', 'load_data', 'export_data', 'get_table_info', 'analyze_data',
+            'insert_document', 'insert_documents', 'find_documents', 'update_document', 'delete_document',
+            'create_vector_index', 'vector_search', 'store_with_embedding', 'text_search', 'get_stats', 'connect_mongodb',
             'execute_command', 'list_processes', 'kill_process', 'get_system_info',
             'evaluate', 'solve_equation', 'convert_units', 'calculate_statistics', 'calculate_financial'
         ]:
             return getattr(module, name)
         if name in ['file_tools', 'pandas_tools', 'wikipedia_tools',
-                   'newspaper_tools', 'arxiv_tools', 'spider_tools', 'duckdb_tools', 'csv_tools', 'json_tools', 'excel_tools', 'xml_tools', 'yaml_tools', 'calculator_tools', 'python_tools', 'shell_tools', 'cot_tools']:
+                   'newspaper_tools', 'arxiv_tools', 'spider_tools', 'duckdb_tools', 'mongodb_tools', 'csv_tools', 'json_tools', 'excel_tools', 'xml_tools', 'yaml_tools', 'calculator_tools', 'python_tools', 'shell_tools', 'cot_tools']:
             return module  # Returns the callable module
         return getattr(module, name)
     else: