tinyagent-py 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl
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.
- tinyagent/code_agent/helper.py +2 -2
- tinyagent/code_agent/modal_sandbox.py +1 -1
- tinyagent/code_agent/providers/base.py +153 -7
- tinyagent/code_agent/providers/modal_provider.py +141 -27
- tinyagent/code_agent/safety.py +6 -2
- tinyagent/code_agent/tiny_code_agent.py +303 -11
- tinyagent/code_agent/utils.py +97 -1
- tinyagent/hooks/__init__.py +3 -1
- tinyagent/hooks/jupyter_notebook_callback.py +1464 -0
- tinyagent/hooks/token_tracker.py +564 -0
- tinyagent/prompts/summarize.yaml +96 -0
- tinyagent/tiny_agent.py +426 -17
- {tinyagent_py-0.0.13.dist-info → tinyagent_py-0.0.15.dist-info}/METADATA +1 -1
- {tinyagent_py-0.0.13.dist-info → tinyagent_py-0.0.15.dist-info}/RECORD +17 -14
- {tinyagent_py-0.0.13.dist-info → tinyagent_py-0.0.15.dist-info}/WHEEL +0 -0
- {tinyagent_py-0.0.13.dist-info → tinyagent_py-0.0.15.dist-info}/licenses/LICENSE +0 -0
- {tinyagent_py-0.0.13.dist-info → tinyagent_py-0.0.15.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,564 @@
|
|
1
|
+
import logging
|
2
|
+
import time
|
3
|
+
from typing import Dict, Any, Optional, List, Union
|
4
|
+
from dataclasses import dataclass, field
|
5
|
+
from collections import defaultdict
|
6
|
+
import json
|
7
|
+
|
8
|
+
@dataclass
|
9
|
+
class UsageStats:
|
10
|
+
"""Represents usage statistics for LLM calls."""
|
11
|
+
prompt_tokens: int = 0
|
12
|
+
completion_tokens: int = 0
|
13
|
+
total_tokens: int = 0
|
14
|
+
cost: float = 0.0
|
15
|
+
call_count: int = 0
|
16
|
+
# Additional fields that LiteLLM might provide
|
17
|
+
thinking_tokens: int = 0
|
18
|
+
reasoning_tokens: int = 0
|
19
|
+
cache_creation_input_tokens: int = 0
|
20
|
+
cache_read_input_tokens: int = 0
|
21
|
+
|
22
|
+
def __add__(self, other: 'UsageStats') -> 'UsageStats':
|
23
|
+
"""Add two UsageStats together."""
|
24
|
+
return UsageStats(
|
25
|
+
prompt_tokens=self.prompt_tokens + other.prompt_tokens,
|
26
|
+
completion_tokens=self.completion_tokens + other.completion_tokens,
|
27
|
+
total_tokens=self.total_tokens + other.total_tokens,
|
28
|
+
cost=self.cost + other.cost,
|
29
|
+
call_count=self.call_count + other.call_count,
|
30
|
+
thinking_tokens=self.thinking_tokens + other.thinking_tokens,
|
31
|
+
reasoning_tokens=self.reasoning_tokens + other.reasoning_tokens,
|
32
|
+
cache_creation_input_tokens=self.cache_creation_input_tokens + other.cache_creation_input_tokens,
|
33
|
+
cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens,
|
34
|
+
)
|
35
|
+
|
36
|
+
def to_dict(self) -> Dict[str, Any]:
|
37
|
+
"""Convert to dictionary."""
|
38
|
+
return {
|
39
|
+
"prompt_tokens": self.prompt_tokens,
|
40
|
+
"completion_tokens": self.completion_tokens,
|
41
|
+
"total_tokens": self.total_tokens,
|
42
|
+
"cost": self.cost,
|
43
|
+
"call_count": self.call_count,
|
44
|
+
"thinking_tokens": self.thinking_tokens,
|
45
|
+
"reasoning_tokens": self.reasoning_tokens,
|
46
|
+
"cache_creation_input_tokens": self.cache_creation_input_tokens,
|
47
|
+
"cache_read_input_tokens": self.cache_read_input_tokens,
|
48
|
+
}
|
49
|
+
|
50
|
+
@classmethod
|
51
|
+
def from_dict(cls, data: Dict[str, Any]) -> 'UsageStats':
|
52
|
+
"""Create from dictionary."""
|
53
|
+
return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
|
54
|
+
|
55
|
+
class TokenTracker:
|
56
|
+
"""
|
57
|
+
A comprehensive token and cost tracker that integrates with TinyAgent's hook system.
|
58
|
+
|
59
|
+
Features:
|
60
|
+
- Accurate tracking using LiteLLM's usage data
|
61
|
+
- Hierarchical tracking for agents with sub-agents
|
62
|
+
- Per-model and per-provider breakdown
|
63
|
+
- Real-time cost calculation
|
64
|
+
- Hook-based integration with TinyAgent
|
65
|
+
"""
|
66
|
+
|
67
|
+
def __init__(
|
68
|
+
self,
|
69
|
+
name: str = "default",
|
70
|
+
parent_tracker: Optional['TokenTracker'] = None,
|
71
|
+
logger: Optional[logging.Logger] = None,
|
72
|
+
enable_detailed_logging: bool = True,
|
73
|
+
track_per_model: bool = True,
|
74
|
+
track_per_provider: bool = True
|
75
|
+
):
|
76
|
+
"""
|
77
|
+
Initialize the TokenTracker.
|
78
|
+
|
79
|
+
Args:
|
80
|
+
name: Name identifier for this tracker
|
81
|
+
parent_tracker: Parent tracker for hierarchical tracking
|
82
|
+
logger: Optional logger instance
|
83
|
+
enable_detailed_logging: Whether to log detailed usage information
|
84
|
+
track_per_model: Whether to track usage per model
|
85
|
+
track_per_provider: Whether to track usage per provider
|
86
|
+
"""
|
87
|
+
self.name = name
|
88
|
+
self.parent_tracker = parent_tracker
|
89
|
+
self.logger = logger or logging.getLogger(__name__)
|
90
|
+
self.enable_detailed_logging = enable_detailed_logging
|
91
|
+
self.track_per_model = track_per_model
|
92
|
+
self.track_per_provider = track_per_provider
|
93
|
+
|
94
|
+
# Overall usage statistics
|
95
|
+
self.total_usage = UsageStats()
|
96
|
+
|
97
|
+
# Per-model tracking
|
98
|
+
self.model_usage: Dict[str, UsageStats] = defaultdict(UsageStats)
|
99
|
+
|
100
|
+
# Per-provider tracking (extracted from model names)
|
101
|
+
self.provider_usage: Dict[str, UsageStats] = defaultdict(UsageStats)
|
102
|
+
|
103
|
+
# Child trackers for hierarchical tracking
|
104
|
+
self.child_trackers: List['TokenTracker'] = []
|
105
|
+
|
106
|
+
# Session tracking
|
107
|
+
self.session_start_time = time.time()
|
108
|
+
self.last_call_time: Optional[float] = None
|
109
|
+
|
110
|
+
# Register with parent if provided
|
111
|
+
if self.parent_tracker:
|
112
|
+
self.parent_tracker.add_child_tracker(self)
|
113
|
+
|
114
|
+
def add_child_tracker(self, child_tracker: 'TokenTracker') -> None:
|
115
|
+
"""Add a child tracker for hierarchical tracking."""
|
116
|
+
if child_tracker not in self.child_trackers:
|
117
|
+
self.child_trackers.append(child_tracker)
|
118
|
+
self.logger.debug(f"Added child tracker '{child_tracker.name}' to '{self.name}'")
|
119
|
+
|
120
|
+
def remove_child_tracker(self, child_tracker: 'TokenTracker') -> None:
|
121
|
+
"""Remove a child tracker."""
|
122
|
+
if child_tracker in self.child_trackers:
|
123
|
+
self.child_trackers.remove(child_tracker)
|
124
|
+
self.logger.debug(f"Removed child tracker '{child_tracker.name}' from '{self.name}'")
|
125
|
+
|
126
|
+
def _extract_provider_from_model(self, model: str) -> str:
|
127
|
+
"""Extract provider name from model string."""
|
128
|
+
# Handle common provider prefixes
|
129
|
+
if "/" in model:
|
130
|
+
return model.split("/")[0]
|
131
|
+
elif model.startswith(("gpt-", "o1", "o3", "o4")):
|
132
|
+
return "openai"
|
133
|
+
elif model.startswith(("claude-", "anthropic/")):
|
134
|
+
return "anthropic"
|
135
|
+
elif model.startswith(("gemini-", "google/")):
|
136
|
+
return "google"
|
137
|
+
elif model.startswith("cohere/"):
|
138
|
+
return "cohere"
|
139
|
+
else:
|
140
|
+
return "unknown"
|
141
|
+
|
142
|
+
def _extract_usage_from_response(self, response: Any) -> Dict[str, Any]:
|
143
|
+
"""Extract usage data from LiteLLM response."""
|
144
|
+
usage_data = {}
|
145
|
+
|
146
|
+
if not response or not hasattr(response, 'usage'):
|
147
|
+
return usage_data
|
148
|
+
|
149
|
+
usage = response.usage
|
150
|
+
|
151
|
+
# Handle both dict and object usage formats
|
152
|
+
if isinstance(usage, dict):
|
153
|
+
usage_data.update(usage)
|
154
|
+
else:
|
155
|
+
# Convert object to dict
|
156
|
+
for attr in dir(usage):
|
157
|
+
if not attr.startswith('_'):
|
158
|
+
value = getattr(usage, attr)
|
159
|
+
if isinstance(value, (int, float)):
|
160
|
+
usage_data[attr] = value
|
161
|
+
|
162
|
+
# Extract cost from LiteLLM response (multiple methods)
|
163
|
+
cost = 0.0
|
164
|
+
|
165
|
+
# Method 1: Check response._hidden_params["response_cost"]
|
166
|
+
try:
|
167
|
+
if hasattr(response, '_hidden_params') and isinstance(response._hidden_params, dict):
|
168
|
+
cost = response._hidden_params.get("response_cost", 0.0)
|
169
|
+
if cost > 0:
|
170
|
+
self.logger.debug(f"Found cost in _hidden_params: ${cost:.6f}")
|
171
|
+
except Exception as e:
|
172
|
+
self.logger.debug(f"Could not extract cost from _hidden_params: {e}")
|
173
|
+
|
174
|
+
# Method 2: Try litellm.completion_cost() as fallback
|
175
|
+
if cost == 0.0:
|
176
|
+
try:
|
177
|
+
import litellm
|
178
|
+
if hasattr(litellm, 'completion_cost'):
|
179
|
+
cost = litellm.completion_cost(completion_response=response)
|
180
|
+
if cost > 0:
|
181
|
+
self.logger.debug(f"Calculated cost using litellm.completion_cost: ${cost:.6f}")
|
182
|
+
except Exception as e:
|
183
|
+
self.logger.debug(f"Could not calculate cost using litellm.completion_cost: {e}")
|
184
|
+
|
185
|
+
# Method 3: Check if cost is already in usage data
|
186
|
+
if cost == 0.0 and 'cost' in usage_data:
|
187
|
+
cost = usage_data.get('cost', 0.0)
|
188
|
+
if cost > 0:
|
189
|
+
self.logger.debug(f"Found cost in usage data: ${cost:.6f}")
|
190
|
+
|
191
|
+
# Add the cost to usage_data
|
192
|
+
usage_data['cost'] = cost
|
193
|
+
|
194
|
+
return usage_data
|
195
|
+
|
196
|
+
def track_llm_call(
|
197
|
+
self,
|
198
|
+
model: str,
|
199
|
+
response: Any,
|
200
|
+
**kwargs
|
201
|
+
) -> None:
|
202
|
+
"""
|
203
|
+
Track a single LLM call using LiteLLM response data.
|
204
|
+
|
205
|
+
Args:
|
206
|
+
model: The model name used
|
207
|
+
response: LiteLLM response object
|
208
|
+
**kwargs: Additional context data
|
209
|
+
"""
|
210
|
+
self.last_call_time = time.time()
|
211
|
+
|
212
|
+
# Extract usage data from LiteLLM response
|
213
|
+
usage_data = self._extract_usage_from_response(response)
|
214
|
+
|
215
|
+
if not usage_data:
|
216
|
+
self.logger.warning(f"No usage data found in response for model {model}")
|
217
|
+
return
|
218
|
+
|
219
|
+
# Create usage stats from response data
|
220
|
+
call_usage = UsageStats(
|
221
|
+
prompt_tokens=usage_data.get('prompt_tokens', 0),
|
222
|
+
completion_tokens=usage_data.get('completion_tokens', 0),
|
223
|
+
total_tokens=usage_data.get('total_tokens', 0),
|
224
|
+
cost=usage_data.get('cost', 0.0),
|
225
|
+
call_count=1,
|
226
|
+
thinking_tokens=usage_data.get('thinking_tokens', 0),
|
227
|
+
reasoning_tokens=usage_data.get('reasoning_tokens', 0),
|
228
|
+
cache_creation_input_tokens=usage_data.get('cache_creation_input_tokens', 0),
|
229
|
+
cache_read_input_tokens=usage_data.get('cache_read_input_tokens', 0),
|
230
|
+
)
|
231
|
+
|
232
|
+
# Update total usage
|
233
|
+
self.total_usage += call_usage
|
234
|
+
|
235
|
+
# Track per-model usage
|
236
|
+
if self.track_per_model:
|
237
|
+
self.model_usage[model] += call_usage
|
238
|
+
|
239
|
+
# Track per-provider usage
|
240
|
+
if self.track_per_provider:
|
241
|
+
provider = self._extract_provider_from_model(model)
|
242
|
+
self.provider_usage[provider] += call_usage
|
243
|
+
|
244
|
+
# Log detailed information if enabled
|
245
|
+
if self.enable_detailed_logging:
|
246
|
+
self.logger.info(
|
247
|
+
f"TokenTracker '{self.name}': {model} call - "
|
248
|
+
f"Tokens: {call_usage.prompt_tokens}+{call_usage.completion_tokens}={call_usage.total_tokens}, "
|
249
|
+
f"Cost: ${call_usage.cost:.6f}"
|
250
|
+
)
|
251
|
+
|
252
|
+
# Log additional token types if present
|
253
|
+
if call_usage.thinking_tokens > 0:
|
254
|
+
self.logger.info(f" Thinking tokens: {call_usage.thinking_tokens}")
|
255
|
+
if call_usage.reasoning_tokens > 0:
|
256
|
+
self.logger.info(f" Reasoning tokens: {call_usage.reasoning_tokens}")
|
257
|
+
if call_usage.cache_creation_input_tokens > 0:
|
258
|
+
self.logger.info(f" Cache creation tokens: {call_usage.cache_creation_input_tokens}")
|
259
|
+
if call_usage.cache_read_input_tokens > 0:
|
260
|
+
self.logger.info(f" Cache read tokens: {call_usage.cache_read_input_tokens}")
|
261
|
+
|
262
|
+
def get_total_usage(self, include_children: bool = False) -> UsageStats:
|
263
|
+
"""
|
264
|
+
Get total usage statistics.
|
265
|
+
|
266
|
+
Args:
|
267
|
+
include_children: Whether to include usage from child trackers
|
268
|
+
|
269
|
+
Returns:
|
270
|
+
UsageStats object with total usage
|
271
|
+
"""
|
272
|
+
total = UsageStats(
|
273
|
+
prompt_tokens=self.total_usage.prompt_tokens,
|
274
|
+
completion_tokens=self.total_usage.completion_tokens,
|
275
|
+
total_tokens=self.total_usage.total_tokens,
|
276
|
+
cost=self.total_usage.cost,
|
277
|
+
call_count=self.total_usage.call_count,
|
278
|
+
thinking_tokens=self.total_usage.thinking_tokens,
|
279
|
+
reasoning_tokens=self.total_usage.reasoning_tokens,
|
280
|
+
cache_creation_input_tokens=self.total_usage.cache_creation_input_tokens,
|
281
|
+
cache_read_input_tokens=self.total_usage.cache_read_input_tokens,
|
282
|
+
)
|
283
|
+
|
284
|
+
if include_children:
|
285
|
+
for child in self.child_trackers:
|
286
|
+
child_usage = child.get_total_usage(include_children=True)
|
287
|
+
total += child_usage
|
288
|
+
|
289
|
+
return total
|
290
|
+
|
291
|
+
def get_model_breakdown(self, include_children: bool = False) -> Dict[str, UsageStats]:
|
292
|
+
"""Get usage breakdown by model."""
|
293
|
+
breakdown = {model: UsageStats(
|
294
|
+
prompt_tokens=stats.prompt_tokens,
|
295
|
+
completion_tokens=stats.completion_tokens,
|
296
|
+
total_tokens=stats.total_tokens,
|
297
|
+
cost=stats.cost,
|
298
|
+
call_count=stats.call_count,
|
299
|
+
thinking_tokens=stats.thinking_tokens,
|
300
|
+
reasoning_tokens=stats.reasoning_tokens,
|
301
|
+
cache_creation_input_tokens=stats.cache_creation_input_tokens,
|
302
|
+
cache_read_input_tokens=stats.cache_read_input_tokens,
|
303
|
+
) for model, stats in self.model_usage.items()}
|
304
|
+
|
305
|
+
if include_children:
|
306
|
+
for child in self.child_trackers:
|
307
|
+
child_breakdown = child.get_model_breakdown(include_children=True)
|
308
|
+
for model, stats in child_breakdown.items():
|
309
|
+
if model in breakdown:
|
310
|
+
breakdown[model] += stats
|
311
|
+
else:
|
312
|
+
breakdown[model] = stats
|
313
|
+
|
314
|
+
return breakdown
|
315
|
+
|
316
|
+
def get_provider_breakdown(self, include_children: bool = False) -> Dict[str, UsageStats]:
|
317
|
+
"""Get usage breakdown by provider."""
|
318
|
+
breakdown = {provider: UsageStats(
|
319
|
+
prompt_tokens=stats.prompt_tokens,
|
320
|
+
completion_tokens=stats.completion_tokens,
|
321
|
+
total_tokens=stats.total_tokens,
|
322
|
+
cost=stats.cost,
|
323
|
+
call_count=stats.call_count,
|
324
|
+
thinking_tokens=stats.thinking_tokens,
|
325
|
+
reasoning_tokens=stats.reasoning_tokens,
|
326
|
+
cache_creation_input_tokens=stats.cache_creation_input_tokens,
|
327
|
+
cache_read_input_tokens=stats.cache_read_input_tokens,
|
328
|
+
) for provider, stats in self.provider_usage.items()}
|
329
|
+
|
330
|
+
if include_children:
|
331
|
+
for child in self.child_trackers:
|
332
|
+
child_breakdown = child.get_provider_breakdown(include_children=True)
|
333
|
+
for provider, stats in child_breakdown.items():
|
334
|
+
if provider in breakdown:
|
335
|
+
breakdown[provider] += stats
|
336
|
+
else:
|
337
|
+
breakdown[provider] = stats
|
338
|
+
|
339
|
+
return breakdown
|
340
|
+
|
341
|
+
def get_session_duration(self) -> float:
|
342
|
+
"""Get session duration in seconds."""
|
343
|
+
return time.time() - self.session_start_time
|
344
|
+
|
345
|
+
def get_detailed_report(self, include_children: bool = True) -> Dict[str, Any]:
|
346
|
+
"""
|
347
|
+
Generate a detailed usage report.
|
348
|
+
|
349
|
+
Args:
|
350
|
+
include_children: Whether to include child tracker data
|
351
|
+
|
352
|
+
Returns:
|
353
|
+
Dictionary containing comprehensive usage information
|
354
|
+
"""
|
355
|
+
total_usage = self.get_total_usage(include_children=include_children)
|
356
|
+
model_breakdown = self.get_model_breakdown(include_children=include_children)
|
357
|
+
provider_breakdown = self.get_provider_breakdown(include_children=include_children)
|
358
|
+
|
359
|
+
report = {
|
360
|
+
"tracker_name": self.name,
|
361
|
+
"session_duration_seconds": self.get_session_duration(),
|
362
|
+
"total_usage": total_usage.to_dict(),
|
363
|
+
"model_breakdown": {model: stats.to_dict() for model, stats in model_breakdown.items()},
|
364
|
+
"provider_breakdown": {provider: stats.to_dict() for provider, stats in provider_breakdown.items()},
|
365
|
+
"child_trackers": []
|
366
|
+
}
|
367
|
+
|
368
|
+
if include_children:
|
369
|
+
for child in self.child_trackers:
|
370
|
+
child_report = child.get_detailed_report(include_children=True)
|
371
|
+
report["child_trackers"].append(child_report)
|
372
|
+
|
373
|
+
return report
|
374
|
+
|
375
|
+
def print_summary(self, include_children: bool = True, detailed: bool = False) -> None:
|
376
|
+
"""Print a summary of usage statistics."""
|
377
|
+
total_usage = self.get_total_usage(include_children=include_children)
|
378
|
+
|
379
|
+
print(f"\n📊 Token Tracker Summary: '{self.name}'")
|
380
|
+
print("=" * 50)
|
381
|
+
print(f"Total Tokens: {total_usage.total_tokens:,}")
|
382
|
+
print(f" • Prompt: {total_usage.prompt_tokens:,}")
|
383
|
+
print(f" • Completion: {total_usage.completion_tokens:,}")
|
384
|
+
if total_usage.thinking_tokens > 0:
|
385
|
+
print(f" • Thinking: {total_usage.thinking_tokens:,}")
|
386
|
+
if total_usage.reasoning_tokens > 0:
|
387
|
+
print(f" • Reasoning: {total_usage.reasoning_tokens:,}")
|
388
|
+
if total_usage.cache_creation_input_tokens > 0:
|
389
|
+
print(f" • Cache Creation: {total_usage.cache_creation_input_tokens:,}")
|
390
|
+
if total_usage.cache_read_input_tokens > 0:
|
391
|
+
print(f" • Cache Read: {total_usage.cache_read_input_tokens:,}")
|
392
|
+
|
393
|
+
print(f"Total Cost: ${total_usage.cost:.6f}")
|
394
|
+
print(f"API Calls: {total_usage.call_count}")
|
395
|
+
print(f"Session Duration: {self.get_session_duration():.1f}s")
|
396
|
+
|
397
|
+
if detailed:
|
398
|
+
model_breakdown = self.get_model_breakdown(include_children=include_children)
|
399
|
+
if model_breakdown:
|
400
|
+
print(f"\n📈 Model Breakdown:")
|
401
|
+
for model, stats in sorted(model_breakdown.items(), key=lambda x: x[1].cost, reverse=True):
|
402
|
+
print(f" {model}: {stats.total_tokens:,} tokens, ${stats.cost:.6f}, {stats.call_count} calls")
|
403
|
+
|
404
|
+
provider_breakdown = self.get_provider_breakdown(include_children=include_children)
|
405
|
+
if provider_breakdown:
|
406
|
+
print(f"\n🏢 Provider Breakdown:")
|
407
|
+
for provider, stats in sorted(provider_breakdown.items(), key=lambda x: x[1].cost, reverse=True):
|
408
|
+
print(f" {provider}: {stats.total_tokens:,} tokens, ${stats.cost:.6f}, {stats.call_count} calls")
|
409
|
+
|
410
|
+
if include_children and self.child_trackers:
|
411
|
+
print(f"\n👥 Child Trackers: {len(self.child_trackers)}")
|
412
|
+
for child in self.child_trackers:
|
413
|
+
child_usage = child.get_total_usage(include_children=True)
|
414
|
+
print(f" • {child.name}: {child_usage.total_tokens:,} tokens, ${child_usage.cost:.6f}")
|
415
|
+
|
416
|
+
def reset_stats(self, reset_children: bool = False) -> None:
|
417
|
+
"""Reset all statistics."""
|
418
|
+
self.total_usage = UsageStats()
|
419
|
+
self.model_usage.clear()
|
420
|
+
self.provider_usage.clear()
|
421
|
+
self.session_start_time = time.time()
|
422
|
+
self.last_call_time = None
|
423
|
+
|
424
|
+
if reset_children:
|
425
|
+
for child in self.child_trackers:
|
426
|
+
child.reset_stats(reset_children=True)
|
427
|
+
|
428
|
+
self.logger.info(f"Reset statistics for tracker '{self.name}'")
|
429
|
+
|
430
|
+
def export_to_json(self, include_children: bool = True) -> str:
|
431
|
+
"""Export tracker data to JSON string."""
|
432
|
+
report = self.get_detailed_report(include_children=include_children)
|
433
|
+
return json.dumps(report, indent=2)
|
434
|
+
|
435
|
+
def save_to_file(self, filepath: str, include_children: bool = True) -> None:
|
436
|
+
"""Save tracker data to a JSON file."""
|
437
|
+
report = self.get_detailed_report(include_children=include_children)
|
438
|
+
with open(filepath, 'w') as f:
|
439
|
+
json.dump(report, f, indent=2)
|
440
|
+
self.logger.info(f"Saved tracker report to {filepath}")
|
441
|
+
|
442
|
+
# Hook methods for TinyAgent integration
|
443
|
+
async def __call__(self, event_name: str, agent: Any, **kwargs) -> None:
|
444
|
+
"""
|
445
|
+
Main hook method that integrates with TinyAgent's callback system.
|
446
|
+
|
447
|
+
Args:
|
448
|
+
event_name: The event name from TinyAgent
|
449
|
+
agent: The TinyAgent instance
|
450
|
+
**kwargs: Event-specific data
|
451
|
+
"""
|
452
|
+
if event_name == "llm_end":
|
453
|
+
response = kwargs.get("response")
|
454
|
+
if response:
|
455
|
+
# Extract model from agent or response
|
456
|
+
model = getattr(agent, 'model', 'unknown')
|
457
|
+
|
458
|
+
# Remove 'response' from kwargs to avoid duplicate argument error
|
459
|
+
filtered_kwargs = {k: v for k, v in kwargs.items() if k != 'response'}
|
460
|
+
self.track_llm_call(model, response, **filtered_kwargs)
|
461
|
+
|
462
|
+
elif event_name == "agent_start":
|
463
|
+
self.logger.debug(f"Agent '{self.name}' started new conversation")
|
464
|
+
|
465
|
+
elif event_name == "agent_end":
|
466
|
+
if self.enable_detailed_logging:
|
467
|
+
total_usage = self.get_total_usage()
|
468
|
+
self.logger.info(
|
469
|
+
f"Agent '{self.name}' completed - "
|
470
|
+
f"Total: {total_usage.total_tokens} tokens, ${total_usage.cost:.6f}"
|
471
|
+
)
|
472
|
+
|
473
|
+
def create_token_tracker(
|
474
|
+
name: str = "main",
|
475
|
+
parent_tracker: Optional[TokenTracker] = None,
|
476
|
+
logger: Optional[logging.Logger] = None,
|
477
|
+
**kwargs
|
478
|
+
) -> TokenTracker:
|
479
|
+
"""
|
480
|
+
Convenience function to create a TokenTracker instance.
|
481
|
+
|
482
|
+
Args:
|
483
|
+
name: Name for the tracker
|
484
|
+
parent_tracker: Parent tracker for hierarchical tracking
|
485
|
+
logger: Logger instance
|
486
|
+
**kwargs: Additional arguments for TokenTracker
|
487
|
+
|
488
|
+
Returns:
|
489
|
+
TokenTracker instance
|
490
|
+
"""
|
491
|
+
return TokenTracker(
|
492
|
+
name=name,
|
493
|
+
parent_tracker=parent_tracker,
|
494
|
+
logger=logger,
|
495
|
+
**kwargs
|
496
|
+
)
|
497
|
+
|
498
|
+
# Example usage
|
499
|
+
async def run_example():
|
500
|
+
"""Example usage of TokenTracker with TinyAgent."""
|
501
|
+
import sys
|
502
|
+
from tinyagent import TinyAgent
|
503
|
+
from tinyagent.hooks.logging_manager import LoggingManager
|
504
|
+
import os
|
505
|
+
|
506
|
+
# Set up logging
|
507
|
+
log_manager = LoggingManager(default_level=logging.INFO)
|
508
|
+
console_handler = logging.StreamHandler(sys.stdout)
|
509
|
+
log_manager.configure_handler(
|
510
|
+
console_handler,
|
511
|
+
format_string='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
512
|
+
level=logging.INFO
|
513
|
+
)
|
514
|
+
|
515
|
+
# Create main token tracker
|
516
|
+
main_tracker = create_token_tracker(
|
517
|
+
name="main_agent",
|
518
|
+
logger=log_manager.get_logger('token_tracker.main'),
|
519
|
+
enable_detailed_logging=True
|
520
|
+
)
|
521
|
+
|
522
|
+
# Create child tracker for sub-agent
|
523
|
+
sub_tracker = create_token_tracker(
|
524
|
+
name="sub_agent",
|
525
|
+
parent_tracker=main_tracker,
|
526
|
+
logger=log_manager.get_logger('token_tracker.sub'),
|
527
|
+
enable_detailed_logging=True
|
528
|
+
)
|
529
|
+
|
530
|
+
# Create main agent with token tracking
|
531
|
+
main_agent = TinyAgent(
|
532
|
+
model="gpt-4o-mini",
|
533
|
+
api_key=os.environ.get("OPENAI_API_KEY"),
|
534
|
+
logger=log_manager.get_logger('main_agent')
|
535
|
+
)
|
536
|
+
main_agent.add_callback(main_tracker)
|
537
|
+
|
538
|
+
# Create sub-agent with different model
|
539
|
+
sub_agent = TinyAgent(
|
540
|
+
model="claude-3-haiku-20240307",
|
541
|
+
api_key=os.environ.get("ANTHROPIC_API_KEY"),
|
542
|
+
logger=log_manager.get_logger('sub_agent')
|
543
|
+
)
|
544
|
+
sub_agent.add_callback(sub_tracker)
|
545
|
+
|
546
|
+
# Run some tasks
|
547
|
+
await main_agent.run("What is the capital of France?")
|
548
|
+
await sub_agent.run("Explain quantum computing in simple terms.")
|
549
|
+
await main_agent.run("Now tell me about the history of Paris.")
|
550
|
+
|
551
|
+
# Print comprehensive summary
|
552
|
+
main_tracker.print_summary(include_children=True, detailed=True)
|
553
|
+
|
554
|
+
# Export report
|
555
|
+
report_json = main_tracker.export_to_json(include_children=True)
|
556
|
+
print(f"\n📄 JSON Report:\n{report_json}")
|
557
|
+
|
558
|
+
# Clean up
|
559
|
+
await main_agent.close()
|
560
|
+
await sub_agent.close()
|
561
|
+
|
562
|
+
if __name__ == "__main__":
|
563
|
+
import asyncio
|
564
|
+
asyncio.run(run_example())
|
@@ -0,0 +1,96 @@
|
|
1
|
+
user_prompt: |-
|
2
|
+
Your task is to create a detailed summary of the conversation so far, paying close attention to the users explicit requests and your previous actions.\n" +
|
3
|
+
This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
|
4
|
+
|
5
|
+
"Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure youve covered all necessary points. In your analysis process:\n" +
|
6
|
+
|
7
|
+
1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:
|
8
|
+
" - The users explicit requests and intents\n" +
|
9
|
+
" - Your approach to addressing the users requests\n" +
|
10
|
+
- Key decisions, technical concepts and code patterns
|
11
|
+
- Specific details like:
|
12
|
+
- file names
|
13
|
+
- full code snippets
|
14
|
+
- function signatures
|
15
|
+
- file edits
|
16
|
+
- Errors that you ran into and how you fixed them
|
17
|
+
- Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
|
18
|
+
2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.
|
19
|
+
|
20
|
+
Your summary should include the following sections:
|
21
|
+
|
22
|
+
"1. Primary Request and Intent: Capture all of the users explicit requests and intents in detail\n" +
|
23
|
+
2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
|
24
|
+
3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
|
25
|
+
4. Errors and fixes: List all errors that you ran into, and how you fixed them. Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
|
26
|
+
5. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
|
27
|
+
"6. All user messages: List ALL user messages that are not tool results. These are critical for understanding the users feedback and changing intent.\n" +
|
28
|
+
6. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
|
29
|
+
7. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
|
30
|
+
"8. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the users explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.\n" +
|
31
|
+
" If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure theres no drift in task interpretation.\n" +
|
32
|
+
|
33
|
+
"Heres an example of how your output should be structured:\n" +
|
34
|
+
|
35
|
+
<example>
|
36
|
+
<analysis>
|
37
|
+
[Your thought process, ensuring all points are covered thoroughly and accurately]
|
38
|
+
</analysis>
|
39
|
+
|
40
|
+
<summary>
|
41
|
+
1. Primary Request and Intent:
|
42
|
+
[Detailed description]
|
43
|
+
|
44
|
+
2. Key Technical Concepts:
|
45
|
+
- [Concept 1]
|
46
|
+
- [Concept 2]
|
47
|
+
- [...]
|
48
|
+
|
49
|
+
3. Files and Code Sections:
|
50
|
+
- [File Name 1]
|
51
|
+
- [Summary of why this file is important]
|
52
|
+
- [Summary of the changes made to this file, if any]
|
53
|
+
- [Important Code Snippet]
|
54
|
+
- [File Name 2]
|
55
|
+
- [Important Code Snippet]
|
56
|
+
- [...]
|
57
|
+
|
58
|
+
4. Errors and fixes:
|
59
|
+
- [Detailed description of error 1]:
|
60
|
+
- [How you fixed the error]
|
61
|
+
- [User feedback on the error if any]
|
62
|
+
- [...]
|
63
|
+
|
64
|
+
5. Problem Solving:
|
65
|
+
[Description of solved problems and ongoing troubleshooting]
|
66
|
+
|
67
|
+
6. All user messages:
|
68
|
+
- [Detailed non tool use user message]
|
69
|
+
- [...]
|
70
|
+
|
71
|
+
7. Pending Tasks:
|
72
|
+
- [Task 1]
|
73
|
+
- [Task 2]
|
74
|
+
- [...]
|
75
|
+
|
76
|
+
8. Current Work:
|
77
|
+
[Precise description of current work]
|
78
|
+
|
79
|
+
9. Optional Next Step:
|
80
|
+
[Optional Next step to take]
|
81
|
+
|
82
|
+
</summary>
|
83
|
+
</example>
|
84
|
+
|
85
|
+
Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response.
|
86
|
+
|
87
|
+
There may be additional summarization instructions provided in the included context. If so, remember to follow these instructions when creating the above summary. Examples of instructions include:
|
88
|
+
<example>
|
89
|
+
## Compact Instructions
|
90
|
+
When summarizing the conversation focus on typescript code changes and also remember the mistakes you made and how you fixed them.
|
91
|
+
</example>
|
92
|
+
|
93
|
+
<example>
|
94
|
+
# Summary instructions
|
95
|
+
When you are using compact - please focus on test output and code changes. Include file reads verbatim.
|
96
|
+
</example>
|