praisonaiagents 0.0.152__py3-none-any.whl → 0.0.154__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
@@ -380,6 +383,7 @@ class Agent:
                     llm_config['base_url'] = base_url
                     if api_key:
                         llm_config['api_key'] = api_key
+                    llm_config['metrics'] = metrics
                     self.llm_instance = LLM(**llm_config)
                 else:
                     # Create LLM with model string and base_url
@@ -387,7 +391,8 @@ class Agent:
                     self.llm_instance = LLM(
                         model=model_name,
                         base_url=base_url,
-                        api_key=api_key
+                        api_key=api_key,
+                        metrics=metrics
                     )
                 self._using_custom_llm = True
             except ImportError as e:
@@ -403,6 +408,9 @@ class Agent:
                 if api_key and 'api_key' not in llm:
                     llm = llm.copy()
                     llm['api_key'] = api_key
+                # Add metrics parameter
+                llm = llm.copy()
+                llm['metrics'] = metrics
                 self.llm_instance = LLM(**llm)  # Pass all dict items as kwargs
                 self._using_custom_llm = True
             except ImportError as e:
@@ -418,6 +426,7 @@ class Agent:
                 llm_params = {'model': llm}
                 if api_key:
                     llm_params['api_key'] = api_key
+                llm_params['metrics'] = metrics
                 self.llm_instance = LLM(**llm_params)
                 self._using_custom_llm = True
                 
@@ -465,6 +474,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
@@ -253,6 +260,12 @@ class LLM:
         self.max_reflect = extra_settings.get('max_reflect', 3)
         self.min_reflect = extra_settings.get('min_reflect', 1)
         self.reasoning_steps = extra_settings.get('reasoning_steps', False)
+        self.metrics = extra_settings.get('metrics', 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
@@ -941,6 +954,10 @@ class LLM:
                         response_text = resp["choices"][0]["message"]["content"]
                         final_response = resp
                         
+                        # Track token usage
+                        if self.metrics:
+                            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 +1135,10 @@ 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
+                                            if self.metrics:
+                                                self._track_token_usage(final_response, self.model)
                                         
                                         # Execute callbacks and display based on verbose setting
                                         if verbose and not interaction_displayed:
@@ -1264,6 +1285,10 @@ 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
+                                if self.metrics:
+                                    self._track_token_usage(final_response, self.model)
                             
                             # Execute callbacks and display based on verbose setting
                             if verbose and not interaction_displayed:
@@ -2861,6 +2886,58 @@ 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
+        
+        # Note: metrics check moved to call sites for performance
+        # This method should only be called when self.metrics=True
+        
+        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

@@ -5,12 +5,14 @@ over HTTP Stream transport, implementing the Streamable HTTP transport protocol.
 """
 
 import asyncio
+import atexit
 import logging
 import threading
 import inspect
 import json
 import time
 import uuid
+import weakref
 from typing import List, Dict, Any, Optional, Callable, Iterable, Union
 from urllib.parse import urlparse, urljoin
 
@@ -25,6 +27,10 @@ logger = logging.getLogger("mcp-http-stream")
 # Global event loop for async operations
 _event_loop = None
 
+# Global registry of active clients for cleanup
+_active_clients = weakref.WeakSet()
+_cleanup_registered = False
+
 def get_event_loop():
     """Get or create a global event loop."""
     global _event_loop
@@ -34,6 +40,31 @@ def get_event_loop():
     return _event_loop
 
 
+def _cleanup_all_clients():
+    """Clean up all active clients at program exit."""
+    if not _active_clients:
+        return
+    
+    # Create a copy to avoid modification during iteration
+    clients_to_cleanup = list(_active_clients)
+    
+    for client in clients_to_cleanup:
+        try:
+            if hasattr(client, '_force_cleanup'):
+                client._force_cleanup()
+        except Exception:
+            # Ignore exceptions during cleanup
+            pass
+
+
+def _register_cleanup():
+    """Register the cleanup function to run at program exit."""
+    global _cleanup_registered
+    if not _cleanup_registered:
+        atexit.register(_cleanup_all_clients)
+        _cleanup_registered = True
+
+
 class HTTPStreamMCPTool:
     """A wrapper for an MCP tool that can be used with praisonaiagents."""
     
@@ -192,14 +223,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 +247,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:
@@ -375,6 +419,7 @@ class HTTPStreamMCPClient:
         self.session = None
         self.tools = []
         self.transport = None
+        self._closed = False
         
         # Set up logging
         if debug:
@@ -383,6 +428,10 @@ class HTTPStreamMCPClient:
             # Set to WARNING by default to hide INFO messages
             logger.setLevel(logging.WARNING)
         
+        # Register this client for cleanup and setup exit handler
+        _active_clients.add(self)
+        _register_cleanup()
+        
         self._initialize()
         
     def _initialize(self):
@@ -443,6 +492,10 @@ class HTTPStreamMCPClient:
                 timeout=self.timeout
             )
             tools.append(wrapper)
+        
+        # Set up cleanup finalizer now that transport and session are created
+        self._finalizer = weakref.finalize(self, self._static_cleanup, 
+                                         self.transport, self._session_context)
             
         return tools
     
@@ -460,7 +513,101 @@ class HTTPStreamMCPClient:
     
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Async context manager exit."""
-        if self.transport:
-            await self.transport.__aexit__(exc_type, exc_val, exc_tb)
-        if hasattr(self, '_session_context') and self._session_context:
-            await self._session_context.__aexit__(exc_type, exc_val, exc_tb)
+        await self.aclose()
+    
+    async def aclose(self):
+        """Async cleanup method to close all resources."""
+        if self._closed:
+            return
+        
+        self._closed = True
+        
+        try:
+            if hasattr(self, '_session_context') and self._session_context:
+                await self._session_context.__aexit__(None, None, None)
+        except Exception:
+            pass
+        
+        try:
+            if self.transport:
+                await self.transport.__aexit__(None, None, None)
+        except Exception:
+            pass
+    
+    def close(self):
+        """Synchronous cleanup method to close all resources."""
+        if self._closed:
+            return
+            
+        try:
+            # Use the global event loop for non-blocking cleanup
+            loop = get_event_loop()
+            if not loop.is_closed():
+                # Schedule cleanup without blocking - add callback for fallback
+                future = asyncio.run_coroutine_threadsafe(self.aclose(), loop)
+                
+                # Add a completion callback for fallback cleanup if async fails
+                def _cleanup_callback(fut):
+                    try:
+                        fut.result()  # This will raise if aclose() failed
+                    except Exception:
+                        # If async cleanup failed, try force cleanup
+                        try:
+                            self._force_cleanup()
+                        except Exception:
+                            pass
+                
+                future.add_done_callback(_cleanup_callback)
+            else:
+                # Event loop is closed, use force cleanup immediately
+                self._force_cleanup()
+        except Exception:
+            # If async scheduling fails, try force cleanup
+            self._force_cleanup()
+    
+    def _force_cleanup(self):
+        """Force cleanup of resources synchronously (for emergencies)."""
+        if self._closed:
+            return
+            
+        self._closed = True
+        
+        # Force close transport session if it exists
+        try:
+            if self.transport and hasattr(self.transport, '_session') and self.transport._session:
+                session = self.transport._session
+                if not session.closed:
+                    # Force close the aiohttp session
+                    if hasattr(session, '_connector') and session._connector:
+                        try:
+                            # Close connector directly
+                            session._connector.close()
+                        except Exception:
+                            pass
+        except Exception:
+            pass
+    
+    @staticmethod
+    def _static_cleanup(transport, session_context):
+        """Static cleanup method for weakref finalizer."""
+        try:
+            # This is called by weakref finalizer, so we can't do async operations
+            # Just ensure any session is closed if possible
+            if transport and hasattr(transport, '_session') and transport._session:
+                session = transport._session
+                if not session.closed and hasattr(session, '_connector'):
+                    try:
+                        session._connector.close()
+                    except Exception:
+                        pass
+        except Exception:
+            pass
+    
+    def __del__(self):
+        """Cleanup when object is garbage collected."""
+        try:
+            if not self._closed:
+                self._force_cleanup()
+        except Exception:
+            # Never raise exceptions in __del__
+            pass