praisonaiagents 0.0.152__py3-none-any.whl → 0.0.153__py3-none-any.whl

praisonaiagents/__init__.py

@@ -60,6 +60,9 @@ try:
         get_telemetry,
         enable_telemetry,
         disable_telemetry,
+        enable_performance_mode,
+        disable_performance_mode,
+        cleanup_telemetry_resources,
         MinimalTelemetry,
         TelemetryCollector
     )
@@ -80,22 +83,66 @@ except ImportError:
     def disable_telemetry():
         pass
     
+    def enable_performance_mode():
+        pass
+    
+    def disable_performance_mode():
+        pass
+    
+    def cleanup_telemetry_resources():
+        pass
+    
     MinimalTelemetry = None
     TelemetryCollector = None
 
 # Add Agents as an alias for PraisonAIAgents
 Agents = PraisonAIAgents
 
-# Apply telemetry auto-instrumentation after all imports
+# Enable PostHog telemetry by default with actual event posting  
+# PostHog events are posted by default unless explicitly disabled
+# Users can:
+#   - Disable completely: PRAISONAI_DISABLE_TELEMETRY=true (or DO_NOT_TRACK=true)
+#   - Enable performance mode: PRAISONAI_PERFORMANCE_MODE=true (minimal overhead, limited events)
+#   - Enable full telemetry: PRAISONAI_FULL_TELEMETRY=true (detailed tracking)
+#   - Legacy opt-in mode: PRAISONAI_AUTO_INSTRUMENT=true
 if _telemetry_available:
     try:
-        # Only instrument if telemetry is enabled
-        _telemetry = get_telemetry()
-        if _telemetry and _telemetry.enabled:
-            from .telemetry.integration import auto_instrument_all
-            auto_instrument_all(_telemetry)
+        import os
+        
+        # Check for explicit disable (respects DO_NOT_TRACK and other disable flags)
+        telemetry_disabled = any([
+            os.environ.get('PRAISONAI_TELEMETRY_DISABLED', '').lower() in ('true', '1', 'yes'),
+            os.environ.get('PRAISONAI_DISABLE_TELEMETRY', '').lower() in ('true', '1', 'yes'),
+            os.environ.get('DO_NOT_TRACK', '').lower() in ('true', '1', 'yes'),
+        ])
+        
+        # Check for performance mode (minimal overhead with limited events)
+        performance_mode = os.environ.get('PRAISONAI_PERFORMANCE_MODE', '').lower() in ('true', '1', 'yes')
+        
+        # Check for full telemetry mode (more detailed tracking)
+        full_telemetry = os.environ.get('PRAISONAI_FULL_TELEMETRY', '').lower() in ('true', '1', 'yes')
+        
+        # Legacy explicit auto-instrument option
+        explicit_auto_instrument = os.environ.get('PRAISONAI_AUTO_INSTRUMENT', '').lower() in ('true', '1', 'yes')
+        
+        # Enable PostHog by default unless explicitly disabled
+        if not telemetry_disabled:
+            _telemetry = get_telemetry()
+            if _telemetry and _telemetry.enabled:
+                from .telemetry.integration import auto_instrument_all
+                
+                # Default: PostHog telemetry is enabled and events are posted
+                # Performance mode can be explicitly enabled for minimal overhead
+                use_performance_mode = performance_mode and not (full_telemetry or explicit_auto_instrument)
+                auto_instrument_all(_telemetry, performance_mode=use_performance_mode)
+                
+                # Track package import for basic usage analytics
+                try:
+                    _telemetry.track_feature_usage("package_import")
+                except Exception:
+                    pass
     except Exception:
-        # Silently fail if there are any issues
+        # Silently fail if there are any issues - never break user applications
         pass
 
 __all__ = [
@@ -135,6 +182,9 @@ __all__ = [
     'get_telemetry',
     'enable_telemetry',
     'disable_telemetry',
+    'enable_performance_mode',
+    'disable_performance_mode',
+    'cleanup_telemetry_resources',
     'MinimalTelemetry',
     'TelemetryCollector'
 ]

praisonaiagents/agent/agent.py

@@ -207,6 +207,7 @@ class Agent:
         use_system_prompt: Optional[bool] = True,
         markdown: bool = True,
         stream: bool = False,
+        metrics: bool = False,
         self_reflect: bool = False,
         max_reflect: int = 3,
         min_reflect: int = 1,
@@ -283,6 +284,8 @@ class Agent:
                 readability and structure. Defaults to True.
             stream (bool, optional): Enable streaming responses from the language model for real-time
                 output when using Agent.start() method. Defaults to False for backward compatibility.
+            metrics (bool, optional): Enable automatic token usage tracking and display summary
+                when tasks complete. Simplifies token monitoring for cost optimization. Defaults to False.
             self_reflect (bool, optional): Enable self-reflection capabilities where the agent
                 evaluates and improves its own responses. Defaults to False.
             max_reflect (int, optional): Maximum number of self-reflection iterations to prevent
@@ -465,6 +468,7 @@ class Agent:
         self.chat_history = []
         self.markdown = markdown
         self.stream = stream
+        self.metrics = metrics
         self.max_reflect = max_reflect
         self.min_reflect = min_reflect
         self.reflect_prompt = reflect_prompt

praisonaiagents/agents/agents.py

@@ -15,6 +15,12 @@ import asyncio
 import uuid
 from enum import Enum
 
+# Import token tracking
+try:
+    from ..telemetry.token_collector import _token_collector
+except ImportError:
+    _token_collector = None
+
 # Task status constants
 class TaskStatus(Enum):
     """Enumeration for task status values to ensure consistency"""
@@ -305,6 +311,11 @@ class PraisonAIAgents:
             task.status = "in progress"
 
         executor_agent = task.agent
+        
+        # Set current agent for token tracking
+        llm = getattr(executor_agent, 'llm', None) or getattr(executor_agent, 'llm_instance', None)
+        if llm and hasattr(llm, 'set_current_agent'):
+            llm.set_current_agent(executor_agent.name)
 
         # Ensure tools are available from both task and agent
         tools = task.tools or []
@@ -401,6 +412,12 @@ Context:
                 agent=executor_agent.name,
                 output_format="RAW"
             )
+            
+            # Add token metrics if available
+            if llm and hasattr(llm, 'last_token_metrics'):
+                token_metrics = llm.last_token_metrics
+                if token_metrics:
+                    task_output.token_metrics = token_metrics
 
             if task.output_json:
                 cleaned = self.clean_json_output(agent_output)
@@ -633,6 +650,11 @@ Context:
             task.status = "in progress"
 
         executor_agent = task.agent
+        
+        # Set current agent for token tracking
+        llm = getattr(executor_agent, 'llm', None) or getattr(executor_agent, 'llm_instance', None)
+        if llm and hasattr(llm, 'set_current_agent'):
+            llm.set_current_agent(executor_agent.name)
 
         task_prompt = f"""
 You need to do the following task: {task.description}.
@@ -749,6 +771,12 @@ Context:
                 agent=executor_agent.name,
                 output_format="RAW"
             )
+            
+            # Add token metrics if available
+            if llm and hasattr(llm, 'last_token_metrics'):
+                token_metrics = llm.last_token_metrics
+                if token_metrics:
+                    task_output.token_metrics = token_metrics
 
             if task.output_json:
                 cleaned = self.clean_json_output(agent_output)
@@ -905,6 +933,18 @@ Context:
         # Run tasks as before
         self.run_all_tasks()
         
+        # Auto-display token metrics if any agent has metrics=True
+        metrics_enabled = any(getattr(agent, 'metrics', False) for agent in self.agents)
+        if metrics_enabled:
+            try:
+                self.display_token_usage()
+            except (ImportError, AttributeError) as e:
+                # Token tracking not available or not properly configured
+                logging.debug(f"Could not auto-display token usage: {e}")
+            except Exception as e:
+                # Log unexpected errors for debugging
+                logging.debug(f"Unexpected error in token metrics display: {e}")
+        
         # Get results
         results = {
             "task_status": self.get_all_tasks_status(),
@@ -924,6 +964,10 @@ Context:
         # Return full results dict if return_dict is True or if no final result was found
         return results
 
+    def run(self, content=None, return_dict=False, **kwargs):
+        """Alias for start() method to provide consistent API with Agent class"""
+        return self.start(content=content, return_dict=return_dict, **kwargs)
+
     def set_state(self, key: str, value: Any) -> None:
         """Set a state value"""
         self._state[key] = value
@@ -1038,6 +1082,77 @@ Context:
                     return True
         
         return False
+
+    def get_token_usage_summary(self) -> Dict[str, Any]:
+        """Get a summary of token usage across all agents and tasks."""
+        if not _token_collector:
+            return {"error": "Token tracking not available"}
+        
+        return _token_collector.get_session_summary()
+    
+    def get_detailed_token_report(self) -> Dict[str, Any]:
+        """Get a detailed token usage report."""
+        if not _token_collector:
+            return {"error": "Token tracking not available"}
+        
+        summary = _token_collector.get_session_summary()
+        recent = _token_collector.get_recent_interactions(limit=20)
+        
+        # Calculate cost estimates (example rates)
+        cost_per_1k_input = 0.0005  # $0.0005 per 1K input tokens
+        cost_per_1k_output = 0.0015  # $0.0015 per 1K output tokens
+        
+        total_metrics = summary.get("total_metrics", {})
+        input_cost = (total_metrics.get("input_tokens", 0) / 1000) * cost_per_1k_input
+        output_cost = (total_metrics.get("output_tokens", 0) / 1000) * cost_per_1k_output
+        total_cost = input_cost + output_cost
+        
+        return {
+            "summary": summary,
+            "recent_interactions": recent,
+            "cost_estimate": {
+                "input_cost": f"${input_cost:.4f}",
+                "output_cost": f"${output_cost:.4f}",
+                "total_cost": f"${total_cost:.4f}",
+                "note": "Cost estimates based on example rates"
+            }
+        }
+    
+    def display_token_usage(self):
+        """Display token usage in a formatted table."""
+        if not _token_collector:
+            print("Token tracking not available")
+            return
+        
+        summary = _token_collector.get_session_summary()
+        
+        print("\n" + "="*50)
+        print("TOKEN USAGE SUMMARY")
+        print("="*50)
+        
+        total_metrics = summary.get("total_metrics", {})
+        print(f"\nTotal Interactions: {summary.get('total_interactions', 0)}")
+        print(f"Total Tokens: {total_metrics.get('total_tokens', 0):,}")
+        print(f"  - Input Tokens: {total_metrics.get('input_tokens', 0):,}")
+        print(f"  - Output Tokens: {total_metrics.get('output_tokens', 0):,}")
+        print(f"  - Cached Tokens: {total_metrics.get('cached_tokens', 0):,}")
+        print(f"  - Reasoning Tokens: {total_metrics.get('reasoning_tokens', 0):,}")
+        
+        # By model
+        by_model = summary.get("by_model", {})
+        if by_model:
+            print("\nUsage by Model:")
+            for model, metrics in by_model.items():
+                print(f"  {model}: {metrics.get('total_tokens', 0):,} tokens")
+        
+        # By agent
+        by_agent = summary.get("by_agent", {})
+        if by_agent:
+            print("\nUsage by Agent:")
+            for agent, metrics in by_agent.items():
+                print(f"  {agent}: {metrics.get('total_tokens', 0):,} tokens")
+        
+        print("="*50 + "\n")
         
     def launch(self, path: str = '/agents', port: int = 8000, host: str = '0.0.0.0', debug: bool = False, protocol: str = "http"):
         """
@@ -1416,4 +1531,4 @@ Context:
             return None
         else:
             display_error(f"Invalid protocol: {protocol}. Choose 'http' or 'mcp'.")
-            return None 
+            return None

praisonaiagents/llm/llm.py

@@ -20,6 +20,13 @@ from ..main import (
 from rich.console import Console
 from rich.live import Live
 
+# Import token tracking
+try:
+    from ..telemetry.token_collector import TokenMetrics, _token_collector
+except ImportError:
+    TokenMetrics = None
+    _token_collector = None
+
 # Logging is already configured in _logging.py via __init__.py
 
 # TODO: Include in-build tool calling in LLM class
@@ -254,6 +261,11 @@ class LLM:
         self.min_reflect = extra_settings.get('min_reflect', 1)
         self.reasoning_steps = extra_settings.get('reasoning_steps', False)
         
+        # Token tracking
+        self.last_token_metrics: Optional[TokenMetrics] = None
+        self.session_token_metrics: Optional[TokenMetrics] = None
+        self.current_agent_name: Optional[str] = None
+        
         # Enable error dropping for cleaner output
         litellm.drop_params = True
         # Enable parameter modification for providers like Anthropic
@@ -941,6 +953,9 @@ class LLM:
                         response_text = resp["choices"][0]["message"]["content"]
                         final_response = resp
                         
+                        # Track token usage
+                        self._track_token_usage(final_response, model)
+                        
                         # Execute callbacks and display based on verbose setting
                         generation_time_val = time.time() - current_time
                         response_content = f"Reasoning:\n{reasoning_content}\n\nAnswer:\n{response_text}" if reasoning_content else response_text
@@ -1118,6 +1133,9 @@ class LLM:
                                             # Handle None content from Gemini
                                             response_content = final_response["choices"][0]["message"].get("content")
                                             response_text = response_content if response_content is not None else ""
+                                            
+                                            # Track token usage
+                                            self._track_token_usage(final_response, self.model)
                                         
                                         # Execute callbacks and display based on verbose setting
                                         if verbose and not interaction_displayed:
@@ -1264,6 +1282,9 @@ class LLM:
                                 # Handle None content from Gemini
                                 response_content = final_response["choices"][0]["message"].get("content")
                                 response_text = response_content if response_content is not None else ""
+                                
+                                # Track token usage
+                                self._track_token_usage(final_response, self.model)
                             
                             # Execute callbacks and display based on verbose setting
                             if verbose and not interaction_displayed:
@@ -2861,6 +2882,55 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                 
         litellm.callbacks = events
 
+    def _track_token_usage(self, response: Dict[str, Any], model: str) -> Optional[TokenMetrics]:
+        """Extract and track token usage from LLM response."""
+        if not TokenMetrics or not _token_collector:
+            return None
+        
+        try:
+            usage = response.get("usage", {})
+            if not usage:
+                return None
+            
+            # Extract token counts
+            metrics = TokenMetrics(
+                input_tokens=usage.get("prompt_tokens", 0),
+                output_tokens=usage.get("completion_tokens", 0),
+                cached_tokens=usage.get("cached_tokens", 0),
+                reasoning_tokens=usage.get("reasoning_tokens", 0),
+                audio_input_tokens=usage.get("audio_input_tokens", 0),
+                audio_output_tokens=usage.get("audio_output_tokens", 0)
+            )
+            
+            # Store metrics
+            self.last_token_metrics = metrics
+            
+            # Update session metrics
+            if not self.session_token_metrics:
+                self.session_token_metrics = TokenMetrics()
+            self.session_token_metrics = self.session_token_metrics + metrics
+            
+            # Track in global collector
+            _token_collector.track_tokens(
+                model=model,
+                agent=self.current_agent_name,
+                metrics=metrics,
+                metadata={
+                    "provider": self.provider,
+                    "stream": False
+                }
+            )
+            
+            return metrics
+            
+        except Exception as e:
+            if self.verbose:
+                logging.warning(f"Failed to track token usage: {e}")
+            return None
+    
+    def set_current_agent(self, agent_name: Optional[str]):
+        """Set the current agent name for token tracking."""
+        self.current_agent_name = agent_name
 
     def _build_completion_params(self, **override_params) -> Dict[str, Any]:
         """Build parameters for litellm completion calls with all necessary config"""

praisonaiagents/main.py

@@ -12,6 +12,12 @@ from rich.markdown import Markdown
 from rich.live import Live
 import asyncio
 
+# Import token metrics if available
+try:
+    from .telemetry.token_collector import TokenMetrics
+except ImportError:
+    TokenMetrics = None
+
 # Logging is already configured in _logging.py via __init__.py
 
 # Global list to store error logs
@@ -415,6 +421,7 @@ class TaskOutput(BaseModel):
     json_dict: Optional[Dict[str, Any]] = None
     agent: str
     output_format: Literal["RAW", "JSON", "Pydantic"] = "RAW"
+    token_metrics: Optional['TokenMetrics'] = None  # Add token metrics field
 
     def json(self) -> Optional[str]:
         if self.output_format == "JSON" and self.json_dict:

praisonaiagents/mcp/mcp_http_stream.py

@@ -192,14 +192,20 @@ class HTTPStreamTransport:
         self._message_queue = asyncio.Queue()
         self._pending_requests = {}
         self._closing = False
+        self._closed = False
         
     async def __aenter__(self):
         self._session = aiohttp.ClientSession()
         return self
     
     async def __aexit__(self, exc_type, exc_val, exc_tb):
+        # Prevent double closing
+        if self._closed:
+            return
+        
         # Set closing flag to stop listener gracefully
         self._closing = True
+        self._closed = True
         
         if self._sse_task:
             self._sse_task.cancel()
@@ -210,6 +216,13 @@ class HTTPStreamTransport:
         if self._session:
             await self._session.close()
     
+    def __del__(self):
+        """Lightweight cleanup during garbage collection."""
+        # Note: We cannot safely run async cleanup in __del__
+        # The best practice is to use async context managers or explicit close() calls
+        pass
+    
+    
     async def send_request(self, request: Dict[str, Any]) -> Union[Dict[str, Any], None]:
         """Send a request to the HTTP Stream endpoint."""
         if not self._session:
@@ -460,7 +473,34 @@ class HTTPStreamMCPClient:
     
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Async context manager exit."""
+        await self.aclose()
+    
+    async def aclose(self):
+        """Async cleanup method to close all resources."""
         if self.transport:
-            await self.transport.__aexit__(exc_type, exc_val, exc_tb)
+            await self.transport.__aexit__(None, None, None)
         if hasattr(self, '_session_context') and self._session_context:
-            await self._session_context.__aexit__(exc_type, exc_val, exc_tb)
+            await self._session_context.__aexit__(None, None, None)
+    
+    def close(self):
+        """Synchronous cleanup method to close all resources."""
+        if hasattr(self, 'transport') and self.transport and not getattr(self.transport, '_closed', False):
+            try:
+                # Use the global event loop for non-blocking cleanup
+                loop = get_event_loop()
+                if not loop.is_closed():
+                    # Schedule cleanup without blocking
+                    asyncio.run_coroutine_threadsafe(self.aclose(), loop)
+            except Exception:
+                # Silently ignore cleanup errors to avoid impacting performance
+                pass
+    
+    def __del__(self):
+        """Cleanup when object is garbage collected."""
+        try:
+            # Simple, lightweight cleanup
+            if hasattr(self, 'transport') and self.transport:
+                self.close()
+        except Exception:
+            # Never raise exceptions in __del__
+            pass

praisonaiagents/telemetry/__init__.py

@@ -5,19 +5,24 @@ This module provides:
 1. Anonymous usage tracking with privacy-first design  
 2. User-friendly performance monitoring and analysis tools
 
-Telemetry is opt-out and can be disabled via environment variables:
+Telemetry can be disabled via environment variables:
 - PRAISONAI_TELEMETRY_DISABLED=true
 - PRAISONAI_DISABLE_TELEMETRY=true  
 - DO_NOT_TRACK=true
 
+Performance monitoring can be optimized via environment variables:
+- PRAISONAI_PERFORMANCE_DISABLED=true (disables performance monitoring overhead)
+- PRAISONAI_FLOW_ANALYSIS_ENABLED=true (enables expensive flow analysis - opt-in only)
+
 No personal data, prompts, or responses are collected.
 
 Performance Monitoring Features:
 - Function performance tracking with detailed statistics
 - API call monitoring and analysis
-- Function execution flow visualization  
+- Function execution flow visualization (opt-in)
 - Performance bottleneck identification
 - Real-time performance reporting
+- External APM metrics export (DataDog, New Relic compatible)
 - CLI interface for easy access
 """
 
@@ -37,7 +42,7 @@ try:
         PerformanceMonitor, performance_monitor,
         monitor_function, track_api_call, get_performance_report,
         get_function_stats, get_api_stats, get_slowest_functions, 
-        get_slowest_apis, clear_performance_data
+        get_slowest_apis, clear_performance_data, export_external_apm_metrics
     )
     from .performance_utils import (
         FunctionFlowAnalyzer, PerformanceAnalyzer,
@@ -58,6 +63,10 @@ __all__ = [
     'force_shutdown_telemetry',
     'MinimalTelemetry',
     'TelemetryCollector',  # For backward compatibility
+    # Performance optimizations
+    'enable_performance_mode',
+    'disable_performance_mode',
+    'cleanup_telemetry_resources',
 ]
 
 # Add performance monitoring to __all__ if available
@@ -81,6 +90,7 @@ if PERFORMANCE_MONITORING_AVAILABLE:
         'get_slowest_functions',
         'get_slowest_apis',
         'clear_performance_data',
+        'export_external_apm_metrics',
         'analyze_function_flow',
         'visualize_execution_flow',
         'analyze_performance_trends',
@@ -114,6 +124,24 @@ def force_shutdown_telemetry():
     _force_shutdown_telemetry()
 
 
+def enable_performance_mode():
+    """Enable performance mode for minimal telemetry overhead."""
+    from .integration import enable_performance_mode as _enable_performance_mode
+    _enable_performance_mode()
+
+
+def disable_performance_mode():
+    """Disable performance mode to resume full telemetry tracking."""
+    from .integration import disable_performance_mode as _disable_performance_mode
+    _disable_performance_mode()
+
+
+def cleanup_telemetry_resources():
+    """Clean up telemetry resources including thread pools and queues."""
+    from .integration import cleanup_telemetry_resources as _cleanup_telemetry_resources
+    _cleanup_telemetry_resources()
+
+
 # Auto-instrumentation and cleanup setup
 _initialized = False
 _atexit_registered = False