praisonaiagents 0.0.152__tar.gz → 0.0.154__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/PKG-INFO +1 -1
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/__init__.py +57 -7
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agent/agent.py +11 -1
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agents/agents.py +116 -1
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/llm/llm.py +77 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/main.py +7 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/mcp/mcp_http_stream.py +151 -4
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/telemetry/__init__.py +31 -3
- praisonaiagents-0.0.154/praisonaiagents/telemetry/integration.py +611 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/telemetry/performance_monitor.py +162 -1
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/telemetry/performance_utils.py +35 -7
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/telemetry/telemetry.py +145 -42
- praisonaiagents-0.0.154/praisonaiagents/telemetry/token_collector.py +170 -0
- praisonaiagents-0.0.154/praisonaiagents/telemetry/token_telemetry.py +89 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents.egg-info/PKG-INFO +1 -1
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents.egg-info/SOURCES.txt +3 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/pyproject.toml +1 -1
- praisonaiagents-0.0.154/tests/test_basic_agents_demo.py +51 -0
- praisonaiagents-0.0.152/praisonaiagents/telemetry/integration.py +0 -310
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/README.md +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/_logging.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/_warning_patch.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agent/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agent/context_agent.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agent/handoff.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agent/image_agent.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agent/router_agent.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agents/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/agents/autoagents.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/approval.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/flow_display.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/guardrails/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/guardrails/guardrail_result.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/guardrails/llm_guardrail.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/knowledge/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/knowledge/chunking.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/knowledge/knowledge.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/llm/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/llm/model_capabilities.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/llm/model_router.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/llm/openai_client.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/mcp/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/mcp/mcp.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/mcp/mcp_sse.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/memory/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/memory/memory.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/process/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/process/process.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/session.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/task/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/task/task.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/telemetry/performance_cli.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/README.md +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/__init__.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/arxiv_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/calculator_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/csv_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/duckdb_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/duckduckgo_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/excel_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/file_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/json_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/mongodb_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/newspaper_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/pandas_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/python_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/searxng_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/shell_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/spider_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/test.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/train/data/generatecot.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/wikipedia_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/xml_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/yaml_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents/tools/yfinance_tools.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents.egg-info/dependency_links.txt +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents.egg-info/requires.txt +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/praisonaiagents.egg-info/top_level.txt +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/setup.cfg +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test-graph-memory.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_context_agent.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_embedding_logging.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_fix_comprehensive.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_gemini_streaming_fix.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_handoff_compatibility.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_http_stream_basic.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_llm_self_reflection_direct.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_ollama_async_fix.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_ollama_fix.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_ollama_sequential_fix.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_posthog_fixed.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_self_reflection_comprehensive.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_self_reflection_fix_simple.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_self_reflection_fix_verification.py +0 -0
- {praisonaiagents-0.0.152 → praisonaiagents-0.0.154}/tests/test_validation_feedback.py +0 -0
@@ -60,6 +60,9 @@ try:
|
|
60
60
|
get_telemetry,
|
61
61
|
enable_telemetry,
|
62
62
|
disable_telemetry,
|
63
|
+
enable_performance_mode,
|
64
|
+
disable_performance_mode,
|
65
|
+
cleanup_telemetry_resources,
|
63
66
|
MinimalTelemetry,
|
64
67
|
TelemetryCollector
|
65
68
|
)
|
@@ -80,22 +83,66 @@ except ImportError:
|
|
80
83
|
def disable_telemetry():
|
81
84
|
pass
|
82
85
|
|
86
|
+
def enable_performance_mode():
|
87
|
+
pass
|
88
|
+
|
89
|
+
def disable_performance_mode():
|
90
|
+
pass
|
91
|
+
|
92
|
+
def cleanup_telemetry_resources():
|
93
|
+
pass
|
94
|
+
|
83
95
|
MinimalTelemetry = None
|
84
96
|
TelemetryCollector = None
|
85
97
|
|
86
98
|
# Add Agents as an alias for PraisonAIAgents
|
87
99
|
Agents = PraisonAIAgents
|
88
100
|
|
89
|
-
#
|
101
|
+
# Enable PostHog telemetry by default with actual event posting
|
102
|
+
# PostHog events are posted by default unless explicitly disabled
|
103
|
+
# Users can:
|
104
|
+
# - Disable completely: PRAISONAI_DISABLE_TELEMETRY=true (or DO_NOT_TRACK=true)
|
105
|
+
# - Enable performance mode: PRAISONAI_PERFORMANCE_MODE=true (minimal overhead, limited events)
|
106
|
+
# - Enable full telemetry: PRAISONAI_FULL_TELEMETRY=true (detailed tracking)
|
107
|
+
# - Legacy opt-in mode: PRAISONAI_AUTO_INSTRUMENT=true
|
90
108
|
if _telemetry_available:
|
91
109
|
try:
|
92
|
-
|
93
|
-
|
94
|
-
|
95
|
-
|
96
|
-
|
110
|
+
import os
|
111
|
+
|
112
|
+
# Check for explicit disable (respects DO_NOT_TRACK and other disable flags)
|
113
|
+
telemetry_disabled = any([
|
114
|
+
os.environ.get('PRAISONAI_TELEMETRY_DISABLED', '').lower() in ('true', '1', 'yes'),
|
115
|
+
os.environ.get('PRAISONAI_DISABLE_TELEMETRY', '').lower() in ('true', '1', 'yes'),
|
116
|
+
os.environ.get('DO_NOT_TRACK', '').lower() in ('true', '1', 'yes'),
|
117
|
+
])
|
118
|
+
|
119
|
+
# Check for performance mode (minimal overhead with limited events)
|
120
|
+
performance_mode = os.environ.get('PRAISONAI_PERFORMANCE_MODE', '').lower() in ('true', '1', 'yes')
|
121
|
+
|
122
|
+
# Check for full telemetry mode (more detailed tracking)
|
123
|
+
full_telemetry = os.environ.get('PRAISONAI_FULL_TELEMETRY', '').lower() in ('true', '1', 'yes')
|
124
|
+
|
125
|
+
# Legacy explicit auto-instrument option
|
126
|
+
explicit_auto_instrument = os.environ.get('PRAISONAI_AUTO_INSTRUMENT', '').lower() in ('true', '1', 'yes')
|
127
|
+
|
128
|
+
# Enable PostHog by default unless explicitly disabled
|
129
|
+
if not telemetry_disabled:
|
130
|
+
_telemetry = get_telemetry()
|
131
|
+
if _telemetry and _telemetry.enabled:
|
132
|
+
from .telemetry.integration import auto_instrument_all
|
133
|
+
|
134
|
+
# Default: PostHog telemetry is enabled and events are posted
|
135
|
+
# Performance mode can be explicitly enabled for minimal overhead
|
136
|
+
use_performance_mode = performance_mode and not (full_telemetry or explicit_auto_instrument)
|
137
|
+
auto_instrument_all(_telemetry, performance_mode=use_performance_mode)
|
138
|
+
|
139
|
+
# Track package import for basic usage analytics
|
140
|
+
try:
|
141
|
+
_telemetry.track_feature_usage("package_import")
|
142
|
+
except Exception:
|
143
|
+
pass
|
97
144
|
except Exception:
|
98
|
-
# Silently fail if there are any issues
|
145
|
+
# Silently fail if there are any issues - never break user applications
|
99
146
|
pass
|
100
147
|
|
101
148
|
__all__ = [
|
@@ -135,6 +182,9 @@ __all__ = [
|
|
135
182
|
'get_telemetry',
|
136
183
|
'enable_telemetry',
|
137
184
|
'disable_telemetry',
|
185
|
+
'enable_performance_mode',
|
186
|
+
'disable_performance_mode',
|
187
|
+
'cleanup_telemetry_resources',
|
138
188
|
'MinimalTelemetry',
|
139
189
|
'TelemetryCollector'
|
140
190
|
]
|
@@ -207,6 +207,7 @@ class Agent:
|
|
207
207
|
use_system_prompt: Optional[bool] = True,
|
208
208
|
markdown: bool = True,
|
209
209
|
stream: bool = False,
|
210
|
+
metrics: bool = False,
|
210
211
|
self_reflect: bool = False,
|
211
212
|
max_reflect: int = 3,
|
212
213
|
min_reflect: int = 1,
|
@@ -283,6 +284,8 @@ class Agent:
|
|
283
284
|
readability and structure. Defaults to True.
|
284
285
|
stream (bool, optional): Enable streaming responses from the language model for real-time
|
285
286
|
output when using Agent.start() method. Defaults to False for backward compatibility.
|
287
|
+
metrics (bool, optional): Enable automatic token usage tracking and display summary
|
288
|
+
when tasks complete. Simplifies token monitoring for cost optimization. Defaults to False.
|
286
289
|
self_reflect (bool, optional): Enable self-reflection capabilities where the agent
|
287
290
|
evaluates and improves its own responses. Defaults to False.
|
288
291
|
max_reflect (int, optional): Maximum number of self-reflection iterations to prevent
|
@@ -380,6 +383,7 @@ class Agent:
|
|
380
383
|
llm_config['base_url'] = base_url
|
381
384
|
if api_key:
|
382
385
|
llm_config['api_key'] = api_key
|
386
|
+
llm_config['metrics'] = metrics
|
383
387
|
self.llm_instance = LLM(**llm_config)
|
384
388
|
else:
|
385
389
|
# Create LLM with model string and base_url
|
@@ -387,7 +391,8 @@ class Agent:
|
|
387
391
|
self.llm_instance = LLM(
|
388
392
|
model=model_name,
|
389
393
|
base_url=base_url,
|
390
|
-
api_key=api_key
|
394
|
+
api_key=api_key,
|
395
|
+
metrics=metrics
|
391
396
|
)
|
392
397
|
self._using_custom_llm = True
|
393
398
|
except ImportError as e:
|
@@ -403,6 +408,9 @@ class Agent:
|
|
403
408
|
if api_key and 'api_key' not in llm:
|
404
409
|
llm = llm.copy()
|
405
410
|
llm['api_key'] = api_key
|
411
|
+
# Add metrics parameter
|
412
|
+
llm = llm.copy()
|
413
|
+
llm['metrics'] = metrics
|
406
414
|
self.llm_instance = LLM(**llm) # Pass all dict items as kwargs
|
407
415
|
self._using_custom_llm = True
|
408
416
|
except ImportError as e:
|
@@ -418,6 +426,7 @@ class Agent:
|
|
418
426
|
llm_params = {'model': llm}
|
419
427
|
if api_key:
|
420
428
|
llm_params['api_key'] = api_key
|
429
|
+
llm_params['metrics'] = metrics
|
421
430
|
self.llm_instance = LLM(**llm_params)
|
422
431
|
self._using_custom_llm = True
|
423
432
|
|
@@ -465,6 +474,7 @@ class Agent:
|
|
465
474
|
self.chat_history = []
|
466
475
|
self.markdown = markdown
|
467
476
|
self.stream = stream
|
477
|
+
self.metrics = metrics
|
468
478
|
self.max_reflect = max_reflect
|
469
479
|
self.min_reflect = min_reflect
|
470
480
|
self.reflect_prompt = reflect_prompt
|
@@ -15,6 +15,12 @@ import asyncio
|
|
15
15
|
import uuid
|
16
16
|
from enum import Enum
|
17
17
|
|
18
|
+
# Import token tracking
|
19
|
+
try:
|
20
|
+
from ..telemetry.token_collector import _token_collector
|
21
|
+
except ImportError:
|
22
|
+
_token_collector = None
|
23
|
+
|
18
24
|
# Task status constants
|
19
25
|
class TaskStatus(Enum):
|
20
26
|
"""Enumeration for task status values to ensure consistency"""
|
@@ -305,6 +311,11 @@ class PraisonAIAgents:
|
|
305
311
|
task.status = "in progress"
|
306
312
|
|
307
313
|
executor_agent = task.agent
|
314
|
+
|
315
|
+
# Set current agent for token tracking
|
316
|
+
llm = getattr(executor_agent, 'llm', None) or getattr(executor_agent, 'llm_instance', None)
|
317
|
+
if llm and hasattr(llm, 'set_current_agent'):
|
318
|
+
llm.set_current_agent(executor_agent.name)
|
308
319
|
|
309
320
|
# Ensure tools are available from both task and agent
|
310
321
|
tools = task.tools or []
|
@@ -401,6 +412,12 @@ Context:
|
|
401
412
|
agent=executor_agent.name,
|
402
413
|
output_format="RAW"
|
403
414
|
)
|
415
|
+
|
416
|
+
# Add token metrics if available
|
417
|
+
if llm and hasattr(llm, 'last_token_metrics'):
|
418
|
+
token_metrics = llm.last_token_metrics
|
419
|
+
if token_metrics:
|
420
|
+
task_output.token_metrics = token_metrics
|
404
421
|
|
405
422
|
if task.output_json:
|
406
423
|
cleaned = self.clean_json_output(agent_output)
|
@@ -633,6 +650,11 @@ Context:
|
|
633
650
|
task.status = "in progress"
|
634
651
|
|
635
652
|
executor_agent = task.agent
|
653
|
+
|
654
|
+
# Set current agent for token tracking
|
655
|
+
llm = getattr(executor_agent, 'llm', None) or getattr(executor_agent, 'llm_instance', None)
|
656
|
+
if llm and hasattr(llm, 'set_current_agent'):
|
657
|
+
llm.set_current_agent(executor_agent.name)
|
636
658
|
|
637
659
|
task_prompt = f"""
|
638
660
|
You need to do the following task: {task.description}.
|
@@ -749,6 +771,12 @@ Context:
|
|
749
771
|
agent=executor_agent.name,
|
750
772
|
output_format="RAW"
|
751
773
|
)
|
774
|
+
|
775
|
+
# Add token metrics if available
|
776
|
+
if llm and hasattr(llm, 'last_token_metrics'):
|
777
|
+
token_metrics = llm.last_token_metrics
|
778
|
+
if token_metrics:
|
779
|
+
task_output.token_metrics = token_metrics
|
752
780
|
|
753
781
|
if task.output_json:
|
754
782
|
cleaned = self.clean_json_output(agent_output)
|
@@ -905,6 +933,18 @@ Context:
|
|
905
933
|
# Run tasks as before
|
906
934
|
self.run_all_tasks()
|
907
935
|
|
936
|
+
# Auto-display token metrics if any agent has metrics=True
|
937
|
+
metrics_enabled = any(getattr(agent, 'metrics', False) for agent in self.agents)
|
938
|
+
if metrics_enabled:
|
939
|
+
try:
|
940
|
+
self.display_token_usage()
|
941
|
+
except (ImportError, AttributeError) as e:
|
942
|
+
# Token tracking not available or not properly configured
|
943
|
+
logging.debug(f"Could not auto-display token usage: {e}")
|
944
|
+
except Exception as e:
|
945
|
+
# Log unexpected errors for debugging
|
946
|
+
logging.debug(f"Unexpected error in token metrics display: {e}")
|
947
|
+
|
908
948
|
# Get results
|
909
949
|
results = {
|
910
950
|
"task_status": self.get_all_tasks_status(),
|
@@ -924,6 +964,10 @@ Context:
|
|
924
964
|
# Return full results dict if return_dict is True or if no final result was found
|
925
965
|
return results
|
926
966
|
|
967
|
+
def run(self, content=None, return_dict=False, **kwargs):
|
968
|
+
"""Alias for start() method to provide consistent API with Agent class"""
|
969
|
+
return self.start(content=content, return_dict=return_dict, **kwargs)
|
970
|
+
|
927
971
|
def set_state(self, key: str, value: Any) -> None:
|
928
972
|
"""Set a state value"""
|
929
973
|
self._state[key] = value
|
@@ -1038,6 +1082,77 @@ Context:
|
|
1038
1082
|
return True
|
1039
1083
|
|
1040
1084
|
return False
|
1085
|
+
|
1086
|
+
def get_token_usage_summary(self) -> Dict[str, Any]:
|
1087
|
+
"""Get a summary of token usage across all agents and tasks."""
|
1088
|
+
if not _token_collector:
|
1089
|
+
return {"error": "Token tracking not available"}
|
1090
|
+
|
1091
|
+
return _token_collector.get_session_summary()
|
1092
|
+
|
1093
|
+
def get_detailed_token_report(self) -> Dict[str, Any]:
|
1094
|
+
"""Get a detailed token usage report."""
|
1095
|
+
if not _token_collector:
|
1096
|
+
return {"error": "Token tracking not available"}
|
1097
|
+
|
1098
|
+
summary = _token_collector.get_session_summary()
|
1099
|
+
recent = _token_collector.get_recent_interactions(limit=20)
|
1100
|
+
|
1101
|
+
# Calculate cost estimates (example rates)
|
1102
|
+
cost_per_1k_input = 0.0005 # $0.0005 per 1K input tokens
|
1103
|
+
cost_per_1k_output = 0.0015 # $0.0015 per 1K output tokens
|
1104
|
+
|
1105
|
+
total_metrics = summary.get("total_metrics", {})
|
1106
|
+
input_cost = (total_metrics.get("input_tokens", 0) / 1000) * cost_per_1k_input
|
1107
|
+
output_cost = (total_metrics.get("output_tokens", 0) / 1000) * cost_per_1k_output
|
1108
|
+
total_cost = input_cost + output_cost
|
1109
|
+
|
1110
|
+
return {
|
1111
|
+
"summary": summary,
|
1112
|
+
"recent_interactions": recent,
|
1113
|
+
"cost_estimate": {
|
1114
|
+
"input_cost": f"${input_cost:.4f}",
|
1115
|
+
"output_cost": f"${output_cost:.4f}",
|
1116
|
+
"total_cost": f"${total_cost:.4f}",
|
1117
|
+
"note": "Cost estimates based on example rates"
|
1118
|
+
}
|
1119
|
+
}
|
1120
|
+
|
1121
|
+
def display_token_usage(self):
|
1122
|
+
"""Display token usage in a formatted table."""
|
1123
|
+
if not _token_collector:
|
1124
|
+
print("Token tracking not available")
|
1125
|
+
return
|
1126
|
+
|
1127
|
+
summary = _token_collector.get_session_summary()
|
1128
|
+
|
1129
|
+
print("\n" + "="*50)
|
1130
|
+
print("TOKEN USAGE SUMMARY")
|
1131
|
+
print("="*50)
|
1132
|
+
|
1133
|
+
total_metrics = summary.get("total_metrics", {})
|
1134
|
+
print(f"\nTotal Interactions: {summary.get('total_interactions', 0)}")
|
1135
|
+
print(f"Total Tokens: {total_metrics.get('total_tokens', 0):,}")
|
1136
|
+
print(f" - Input Tokens: {total_metrics.get('input_tokens', 0):,}")
|
1137
|
+
print(f" - Output Tokens: {total_metrics.get('output_tokens', 0):,}")
|
1138
|
+
print(f" - Cached Tokens: {total_metrics.get('cached_tokens', 0):,}")
|
1139
|
+
print(f" - Reasoning Tokens: {total_metrics.get('reasoning_tokens', 0):,}")
|
1140
|
+
|
1141
|
+
# By model
|
1142
|
+
by_model = summary.get("by_model", {})
|
1143
|
+
if by_model:
|
1144
|
+
print("\nUsage by Model:")
|
1145
|
+
for model, metrics in by_model.items():
|
1146
|
+
print(f" {model}: {metrics.get('total_tokens', 0):,} tokens")
|
1147
|
+
|
1148
|
+
# By agent
|
1149
|
+
by_agent = summary.get("by_agent", {})
|
1150
|
+
if by_agent:
|
1151
|
+
print("\nUsage by Agent:")
|
1152
|
+
for agent, metrics in by_agent.items():
|
1153
|
+
print(f" {agent}: {metrics.get('total_tokens', 0):,} tokens")
|
1154
|
+
|
1155
|
+
print("="*50 + "\n")
|
1041
1156
|
|
1042
1157
|
def launch(self, path: str = '/agents', port: int = 8000, host: str = '0.0.0.0', debug: bool = False, protocol: str = "http"):
|
1043
1158
|
"""
|
@@ -1416,4 +1531,4 @@ Context:
|
|
1416
1531
|
return None
|
1417
1532
|
else:
|
1418
1533
|
display_error(f"Invalid protocol: {protocol}. Choose 'http' or 'mcp'.")
|
1419
|
-
return None
|
1534
|
+
return None
|
@@ -20,6 +20,13 @@ from ..main import (
|
|
20
20
|
from rich.console import Console
|
21
21
|
from rich.live import Live
|
22
22
|
|
23
|
+
# Import token tracking
|
24
|
+
try:
|
25
|
+
from ..telemetry.token_collector import TokenMetrics, _token_collector
|
26
|
+
except ImportError:
|
27
|
+
TokenMetrics = None
|
28
|
+
_token_collector = None
|
29
|
+
|
23
30
|
# Logging is already configured in _logging.py via __init__.py
|
24
31
|
|
25
32
|
# TODO: Include in-build tool calling in LLM class
|
@@ -253,6 +260,12 @@ class LLM:
|
|
253
260
|
self.max_reflect = extra_settings.get('max_reflect', 3)
|
254
261
|
self.min_reflect = extra_settings.get('min_reflect', 1)
|
255
262
|
self.reasoning_steps = extra_settings.get('reasoning_steps', False)
|
263
|
+
self.metrics = extra_settings.get('metrics', False)
|
264
|
+
|
265
|
+
# Token tracking
|
266
|
+
self.last_token_metrics: Optional[TokenMetrics] = None
|
267
|
+
self.session_token_metrics: Optional[TokenMetrics] = None
|
268
|
+
self.current_agent_name: Optional[str] = None
|
256
269
|
|
257
270
|
# Enable error dropping for cleaner output
|
258
271
|
litellm.drop_params = True
|
@@ -941,6 +954,10 @@ class LLM:
|
|
941
954
|
response_text = resp["choices"][0]["message"]["content"]
|
942
955
|
final_response = resp
|
943
956
|
|
957
|
+
# Track token usage
|
958
|
+
if self.metrics:
|
959
|
+
self._track_token_usage(final_response, model)
|
960
|
+
|
944
961
|
# Execute callbacks and display based on verbose setting
|
945
962
|
generation_time_val = time.time() - current_time
|
946
963
|
response_content = f"Reasoning:\n{reasoning_content}\n\nAnswer:\n{response_text}" if reasoning_content else response_text
|
@@ -1118,6 +1135,10 @@ class LLM:
|
|
1118
1135
|
# Handle None content from Gemini
|
1119
1136
|
response_content = final_response["choices"][0]["message"].get("content")
|
1120
1137
|
response_text = response_content if response_content is not None else ""
|
1138
|
+
|
1139
|
+
# Track token usage
|
1140
|
+
if self.metrics:
|
1141
|
+
self._track_token_usage(final_response, self.model)
|
1121
1142
|
|
1122
1143
|
# Execute callbacks and display based on verbose setting
|
1123
1144
|
if verbose and not interaction_displayed:
|
@@ -1264,6 +1285,10 @@ class LLM:
|
|
1264
1285
|
# Handle None content from Gemini
|
1265
1286
|
response_content = final_response["choices"][0]["message"].get("content")
|
1266
1287
|
response_text = response_content if response_content is not None else ""
|
1288
|
+
|
1289
|
+
# Track token usage
|
1290
|
+
if self.metrics:
|
1291
|
+
self._track_token_usage(final_response, self.model)
|
1267
1292
|
|
1268
1293
|
# Execute callbacks and display based on verbose setting
|
1269
1294
|
if verbose and not interaction_displayed:
|
@@ -2861,6 +2886,58 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
|
2861
2886
|
|
2862
2887
|
litellm.callbacks = events
|
2863
2888
|
|
2889
|
+
def _track_token_usage(self, response: Dict[str, Any], model: str) -> Optional[TokenMetrics]:
|
2890
|
+
"""Extract and track token usage from LLM response."""
|
2891
|
+
if not TokenMetrics or not _token_collector:
|
2892
|
+
return None
|
2893
|
+
|
2894
|
+
# Note: metrics check moved to call sites for performance
|
2895
|
+
# This method should only be called when self.metrics=True
|
2896
|
+
|
2897
|
+
try:
|
2898
|
+
usage = response.get("usage", {})
|
2899
|
+
if not usage:
|
2900
|
+
return None
|
2901
|
+
|
2902
|
+
# Extract token counts
|
2903
|
+
metrics = TokenMetrics(
|
2904
|
+
input_tokens=usage.get("prompt_tokens", 0),
|
2905
|
+
output_tokens=usage.get("completion_tokens", 0),
|
2906
|
+
cached_tokens=usage.get("cached_tokens", 0),
|
2907
|
+
reasoning_tokens=usage.get("reasoning_tokens", 0),
|
2908
|
+
audio_input_tokens=usage.get("audio_input_tokens", 0),
|
2909
|
+
audio_output_tokens=usage.get("audio_output_tokens", 0)
|
2910
|
+
)
|
2911
|
+
|
2912
|
+
# Store metrics
|
2913
|
+
self.last_token_metrics = metrics
|
2914
|
+
|
2915
|
+
# Update session metrics
|
2916
|
+
if not self.session_token_metrics:
|
2917
|
+
self.session_token_metrics = TokenMetrics()
|
2918
|
+
self.session_token_metrics = self.session_token_metrics + metrics
|
2919
|
+
|
2920
|
+
# Track in global collector
|
2921
|
+
_token_collector.track_tokens(
|
2922
|
+
model=model,
|
2923
|
+
agent=self.current_agent_name,
|
2924
|
+
metrics=metrics,
|
2925
|
+
metadata={
|
2926
|
+
"provider": self.provider,
|
2927
|
+
"stream": False
|
2928
|
+
}
|
2929
|
+
)
|
2930
|
+
|
2931
|
+
return metrics
|
2932
|
+
|
2933
|
+
except Exception as e:
|
2934
|
+
if self.verbose:
|
2935
|
+
logging.warning(f"Failed to track token usage: {e}")
|
2936
|
+
return None
|
2937
|
+
|
2938
|
+
def set_current_agent(self, agent_name: Optional[str]):
|
2939
|
+
"""Set the current agent name for token tracking."""
|
2940
|
+
self.current_agent_name = agent_name
|
2864
2941
|
|
2865
2942
|
def _build_completion_params(self, **override_params) -> Dict[str, Any]:
|
2866
2943
|
"""Build parameters for litellm completion calls with all necessary config"""
|
@@ -12,6 +12,12 @@ from rich.markdown import Markdown
|
|
12
12
|
from rich.live import Live
|
13
13
|
import asyncio
|
14
14
|
|
15
|
+
# Import token metrics if available
|
16
|
+
try:
|
17
|
+
from .telemetry.token_collector import TokenMetrics
|
18
|
+
except ImportError:
|
19
|
+
TokenMetrics = None
|
20
|
+
|
15
21
|
# Logging is already configured in _logging.py via __init__.py
|
16
22
|
|
17
23
|
# Global list to store error logs
|
@@ -415,6 +421,7 @@ class TaskOutput(BaseModel):
|
|
415
421
|
json_dict: Optional[Dict[str, Any]] = None
|
416
422
|
agent: str
|
417
423
|
output_format: Literal["RAW", "JSON", "Pydantic"] = "RAW"
|
424
|
+
token_metrics: Optional['TokenMetrics'] = None # Add token metrics field
|
418
425
|
|
419
426
|
def json(self) -> Optional[str]:
|
420
427
|
if self.output_format == "JSON" and self.json_dict:
|