ragaai-catalyst 2.1.3b0__py3-none-any.whl → 2.1.4b1__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.
Files changed (24) hide show
  1. ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py +36 -10
  2. ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py +213 -76
  3. ragaai_catalyst/tracers/agentic_tracing/tracers/base.py +568 -107
  4. ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py +325 -0
  5. ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py +207 -81
  6. ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py +208 -58
  7. ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py +2 -0
  8. ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py +125 -28
  9. ragaai_catalyst/tracers/agentic_tracing/tracers/user_interaction_tracer.py +86 -0
  10. ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py +9 -51
  11. ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py +83 -0
  12. ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py +26 -0
  13. ragaai_catalyst/tracers/agentic_tracing/utils/get_user_trace_metrics.py +28 -0
  14. ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py +45 -15
  15. ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json +2476 -2122
  16. ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py +59 -0
  17. ragaai_catalyst/tracers/agentic_tracing/utils/trace_utils.py +23 -0
  18. ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py +284 -15
  19. ragaai_catalyst/tracers/tracer.py +80 -8
  20. ragaai_catalyst-2.1.4b1.dist-info/METADATA +431 -0
  21. {ragaai_catalyst-2.1.3b0.dist-info → ragaai_catalyst-2.1.4b1.dist-info}/RECORD +23 -18
  22. ragaai_catalyst-2.1.3b0.dist-info/METADATA +0 -43
  23. {ragaai_catalyst-2.1.3b0.dist-info → ragaai_catalyst-2.1.4b1.dist-info}/WHEEL +0 -0
  24. {ragaai_catalyst-2.1.3b0.dist-info → ragaai_catalyst-2.1.4b1.dist-info}/top_level.txt +0 -0
@@ -173,7 +173,26 @@ class LLMCall:
173
173
  duration: float = field(default=0)
174
174
 
175
175
  class Component:
176
- def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, data: Dict[str, Any]={}, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
176
+ def __init__(
177
+ self,
178
+ id: str,
179
+ hash_id: str,
180
+ source_hash_id: str,
181
+ type: str,
182
+ name: str,
183
+ start_time: str,
184
+ end_time: str,
185
+ parent_id: int,
186
+ info: Dict[str, Any],
187
+ extra_info: Optional[Dict[str, Any]] = None,
188
+ data: Dict[str, Any]={},
189
+ metadata: Optional[Dict[str, Any]] = None,
190
+ metrics: Optional[List[Dict[str, Any]]] = None,
191
+ feedback: Optional[Any] = None,
192
+ network_calls: Optional[List[NetworkCall]] = None,
193
+ interactions: Optional[List[Union[Interaction, Dict]]] = None,
194
+ error: Optional[Dict[str, Any]] = None):
195
+
177
196
  self.id = id
178
197
  self.hash_id = hash_id
179
198
  self.source_hash_id = source_hash_id
@@ -185,7 +204,9 @@ class Component:
185
204
  self.info = info
186
205
  self.extra_info = extra_info
187
206
  self.data = data
188
- self.error = error
207
+ self.metadata = metadata or {}
208
+ self.metrics = metrics or []
209
+ self.feedback = feedback
189
210
  self.network_calls = network_calls or []
190
211
  self.interactions = []
191
212
  self.error = error
@@ -217,22 +238,25 @@ class Component:
217
238
  "extra_info": self.extra_info,
218
239
  "error": self.error,
219
240
  "data": self.data,
220
- "error": self.error,
241
+ "metadata": self.metadata,
242
+ "metrics": self.metrics,
243
+ "feedback": self.feedback,
221
244
  "network_calls": [call.to_dict() if hasattr(call, 'to_dict') else call for call in self.network_calls],
222
245
  "interactions": self.interactions
223
246
  }
224
247
 
225
248
  class LLMComponent(Component):
226
- def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
227
- super().__init__(id, hash_id, source_hash_id, type, name, start_time, end_time, parent_id, info, extra_info, data, network_calls, interactions, error)
249
+ def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, data: Dict[str, Any]={}, metadata: Optional[Dict[str, Any]] = None, metrics: Optional[List[Dict[str, Any]]] = None, feedback: Optional[Any] = None, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
250
+ super().__init__(id, hash_id, source_hash_id, type, name, start_time, end_time, parent_id, info, extra_info, data, metadata, metrics, feedback, network_calls, interactions, error)
228
251
 
229
252
  class AgentComponent(Component):
230
- def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
231
- super().__init__(id, hash_id, source_hash_id, type, name, start_time, end_time, parent_id, info, extra_info, data, network_calls, interactions, error)
253
+ def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, data: Dict[str, Any]={}, metadata: Optional[Dict[str, Any]] = None, metrics: Optional[List[Dict[str, Any]]] = None, feedback: Optional[Any] = None, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
254
+ super().__init__(id, hash_id, source_hash_id, type, name, start_time, end_time, parent_id, info, extra_info, data, metadata, metrics, feedback, network_calls, interactions, error)
232
255
 
233
256
  class ToolComponent(Component):
234
- def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
235
- super().__init__(id, hash_id, source_hash_id, type, name, start_time, end_time, parent_id, info, extra_info, data, network_calls, interactions, error)
257
+ def __init__(self, id: str, hash_id: str, source_hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], extra_info: Optional[Dict[str, Any]] = None, data: Dict[str, Any]={}, metadata: Optional[Dict[str, Any]] = None, metrics: Optional[List[Dict[str, Any]]] = None, feedback: Optional[Any] = None, network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
258
+ super().__init__(id, hash_id, source_hash_id, type, name, start_time, end_time, parent_id, info, extra_info, data, metadata, metrics, feedback, network_calls, interactions, error)
259
+
236
260
 
237
261
  @dataclass
238
262
  class ComponentInfo:
@@ -247,8 +271,9 @@ class ComponentInfo:
247
271
  cost: Optional[Dict[str, float]] = None
248
272
 
249
273
  class Trace:
250
- def __init__(self, id: str, project_name: str, start_time: str, end_time: str, metadata: Optional[Metadata] = None, data: Optional[List[Dict[str, Any]]] = None, replays: Optional[Dict[str, Any]] = None):
274
+ def __init__(self, id: str, trace_name: str, project_name: str, start_time: str, end_time: str, metadata: Optional[Metadata] = None, data: Optional[List[Dict[str, Any]]] = None, replays: Optional[Dict[str, Any]] = None):
251
275
  self.id = id
276
+ self.trace_name = trace_name
252
277
  self.project_name = project_name
253
278
  self.start_time = start_time
254
279
  self.end_time = end_time
@@ -259,6 +284,7 @@ class Trace:
259
284
  def to_dict(self):
260
285
  return {
261
286
  "id": self.id,
287
+ "trace_name": self.trace_name,
262
288
  "project_name": self.project_name,
263
289
  "start_time": self.start_time,
264
290
  "end_time": self.end_time,
@@ -1,3 +1,4 @@
1
+ import os
1
2
  import functools
2
3
  import uuid
3
4
  from datetime import datetime
@@ -7,6 +8,15 @@ from ..utils.unique_decorator import mydecorator, generate_unique_hash_simple
7
8
  import contextvars
8
9
  import asyncio
9
10
  from ..utils.file_name_tracker import TrackName
11
+ from ..utils.span_attributes import SpanAttributes
12
+ import logging
13
+
14
+ logger = logging.getLogger(__name__)
15
+ logging_level = (
16
+ logger.setLevel(logging.DEBUG)
17
+ if os.getenv("DEBUG")
18
+ else logger.setLevel(logging.INFO)
19
+ )
10
20
 
11
21
 
12
22
  class AgentTracerMixin:
@@ -16,42 +26,89 @@ class AgentTracerMixin:
16
26
  self.current_agent_id = contextvars.ContextVar("agent_id", default=None)
17
27
  self.current_agent_name = contextvars.ContextVar("agent_name", default=None)
18
28
  self.agent_children = contextvars.ContextVar("agent_children", default=[])
19
- self.component_network_calls = contextvars.ContextVar("component_network_calls", default={})
20
- self.component_user_interaction = contextvars.ContextVar("component_user_interaction", default={})
21
- self.version = contextvars.ContextVar("version", default="1.0.0")
29
+ self.component_network_calls = contextvars.ContextVar(
30
+ "component_network_calls", default={}
31
+ )
32
+ self.component_user_interaction = contextvars.ContextVar(
33
+ "component_user_interaction", default={}
34
+ )
35
+ self.version = contextvars.ContextVar("version", default=None)
22
36
  self.agent_type = contextvars.ContextVar("agent_type", default="generic")
23
37
  self.capabilities = contextvars.ContextVar("capabilities", default=[])
24
38
  self.start_time = contextvars.ContextVar("start_time", default=None)
25
39
  self.input_data = contextvars.ContextVar("input_data", default=None)
26
40
  self.gt = None
27
41
 
42
+ self.span_attributes_dict = {}
43
+
44
+ # Add auto instrument flags
45
+ self.auto_instrument_agent = False
46
+ self.auto_instrument_user_interaction = False
47
+ self.auto_instrument_network = False
48
+
49
+ def trace_agent(
50
+ self,
51
+ name: str,
52
+ agent_type: str = None,
53
+ version: str = None,
54
+ capabilities: List[str] = None,
55
+ tags: List[str] = [],
56
+ metadata: Dict[str, Any] = {},
57
+ metrics: List[Dict[str, Any]] = [],
58
+ feedback: Optional[Any] = None,
59
+ ):
60
+ if name not in self.span_attributes_dict:
61
+ self.span_attributes_dict[name] = SpanAttributes(name)
62
+ if tags:
63
+ self.span(name).add_tags(tags)
64
+ if metadata:
65
+ self.span(name).add_metadata(metadata)
66
+ if metrics:
67
+ if isinstance(metrics, dict):
68
+ metrics = [metrics]
69
+ for metric in metrics:
70
+ try:
71
+ self.span(name).add_metrics(
72
+ name=metric["name"],
73
+ score=metric["score"],
74
+ reasoning=metric.get("reasoning", ""),
75
+ cost=metric.get("cost", None),
76
+ latency=metric.get("latency", None),
77
+ metadata=metric.get("metadata", {}),
78
+ config=metric.get("config", {}),
79
+ )
80
+ except KeyError as e:
81
+ logger.error(f"Error adding metric: {e}")
82
+ if feedback:
83
+ self.span(name).add_feedback(feedback)
28
84
 
29
- def trace_agent(self, name: str, agent_type: str = None, version: str = None, capabilities: List[str] = None):
30
85
  def decorator(target):
31
86
  # Check if target is a class
32
87
  is_class = isinstance(target, type)
33
88
  tracer = self # Store reference to tracer instance
34
- top_level_hash_id = generate_unique_hash_simple(target) # Generate hash based on the decorated target code
89
+ top_level_hash_id = generate_unique_hash_simple(
90
+ target
91
+ ) # Generate hash based on the decorated target code
35
92
  self.version.set(version)
36
93
  self.agent_type.set(agent_type)
37
94
  self.capabilities.set(capabilities)
38
-
95
+
39
96
  if is_class:
40
97
  # Store original __init__
41
98
  original_init = target.__init__
42
-
99
+
43
100
  def wrapped_init(self, *args, **kwargs):
44
- self.gt = kwargs.get('gt', None) if kwargs else None
101
+ self.gt = kwargs.get("gt", None) if kwargs else None
45
102
  # Set agent context before initializing
46
103
  component_id = str(uuid.uuid4())
47
104
  hash_id = top_level_hash_id
48
-
105
+
49
106
  # Store the component ID in the instance
50
107
  self._agent_component_id = component_id
51
-
108
+
52
109
  # Get parent agent ID if exists
53
110
  parent_agent_id = tracer.current_agent_id.get()
54
-
111
+
55
112
  # Create agent component
56
113
  agent_component = tracer.create_agent_component(
57
114
  component_id=component_id,
@@ -65,14 +122,14 @@ class AgentTracerMixin:
65
122
  input_data=tracer._sanitize_input(args, kwargs),
66
123
  output_data=None,
67
124
  children=[],
68
- parent_id=parent_agent_id
125
+ parent_id=parent_agent_id,
69
126
  )
70
-
127
+
71
128
  # Store component for later updates
72
- if not hasattr(tracer, '_agent_components'):
129
+ if not hasattr(tracer, "_agent_components"):
73
130
  tracer._agent_components = {}
74
131
  tracer._agent_components[component_id] = agent_component
75
-
132
+
76
133
  # If this is a nested agent, add it to parent's children
77
134
  if parent_agent_id:
78
135
  parent_children = tracer.agent_children.get()
@@ -81,61 +138,82 @@ class AgentTracerMixin:
81
138
  else:
82
139
  # Only add to root components if no parent
83
140
  tracer.add_component(agent_component)
84
-
141
+
85
142
  # Call original __init__ with this agent as current
86
143
  token = tracer.current_agent_id.set(component_id)
87
144
  try:
88
145
  original_init(self, *args, **kwargs)
89
146
  finally:
90
147
  tracer.current_agent_id.reset(token)
91
-
148
+
92
149
  # Wrap all public methods to track execution
93
150
  for attr_name in dir(target):
94
- if not attr_name.startswith('_'):
151
+ if not attr_name.startswith("_"):
95
152
  attr_value = getattr(target, attr_name)
96
153
  if callable(attr_value):
154
+
97
155
  def wrap_method(method):
98
156
  @self.file_tracker.trace_decorator
99
157
  @functools.wraps(method)
100
158
  def wrapped_method(self, *args, **kwargs):
101
- self.gt = kwargs.get('gt', None) if kwargs else None
159
+ self.gt = kwargs.get("gt", None) if kwargs else None
102
160
  # Set this agent as current during method execution
103
- token = tracer.current_agent_id.set(self._agent_component_id)
104
-
161
+ token = tracer.current_agent_id.set(
162
+ self._agent_component_id
163
+ )
164
+
105
165
  # Store parent's children before setting new empty list
106
166
  parent_children = tracer.agent_children.get()
107
167
  children_token = tracer.agent_children.set([])
108
-
168
+
109
169
  try:
110
170
  start_time = datetime.now()
111
171
  result = method(self, *args, **kwargs)
112
-
172
+
113
173
  # Update agent component with method result
114
- if hasattr(tracer, '_agent_components'):
115
- component = tracer._agent_components.get(self._agent_component_id)
174
+ if hasattr(tracer, "_agent_components"):
175
+ component = tracer._agent_components.get(
176
+ self._agent_component_id
177
+ )
116
178
  if component:
117
- component['data']['output'] = tracer._sanitize_output(result)
118
- component['data']['input'] = tracer._sanitize_input(args, kwargs)
119
- component['start_time'] = start_time.isoformat()
120
-
179
+ component["data"]["output"] = (
180
+ tracer._sanitize_output(result)
181
+ )
182
+ component["data"]["input"] = (
183
+ tracer._sanitize_input(args, kwargs)
184
+ )
185
+ component["start_time"] = (
186
+ start_time.isoformat()
187
+ )
188
+
121
189
  # Get children accumulated during method execution
122
190
  children = tracer.agent_children.get()
123
191
  if children:
124
- if 'children' not in component['data']:
125
- component['data']['children'] = []
126
- component['data']['children'].extend(children)
127
-
192
+ if (
193
+ "children"
194
+ not in component["data"]
195
+ ):
196
+ component["data"][
197
+ "children"
198
+ ] = []
199
+ component["data"][
200
+ "children"
201
+ ].extend(children)
202
+
128
203
  # Add this component as a child to parent's children list
129
204
  parent_children.append(component)
130
- tracer.agent_children.set(parent_children)
205
+ tracer.agent_children.set(
206
+ parent_children
207
+ )
131
208
  return result
132
209
  finally:
133
210
  tracer.current_agent_id.reset(token)
134
211
  tracer.agent_children.reset(children_token)
212
+
135
213
  return wrapped_method
136
-
214
+
137
215
  setattr(target, attr_name, wrap_method(attr_value))
138
-
216
+
139
217
  # Replace __init__ with wrapped version
140
218
  target.__init__ = wrapped_init
141
219
  return target
@@ -143,17 +221,40 @@ class AgentTracerMixin:
143
221
  # For function decorators, use existing sync/async tracing
144
222
  is_async = asyncio.iscoroutinefunction(target)
145
223
  if is_async:
224
+
146
225
  async def wrapper(*args, **kwargs):
147
- return await self._trace_agent_execution(target, name, agent_type, version, capabilities, top_level_hash_id, *args, **kwargs)
226
+ return await self._trace_agent_execution(
227
+ target,
228
+ name,
229
+ agent_type,
230
+ version,
231
+ capabilities,
232
+ top_level_hash_id,
233
+ *args,
234
+ **kwargs,
235
+ )
236
+
148
237
  return wrapper
149
238
  else:
239
+
150
240
  def wrapper(*args, **kwargs):
151
- return self._trace_sync_agent_execution(target, name, agent_type, version, capabilities, *args, **kwargs)
241
+ return self._trace_sync_agent_execution(
242
+ target,
243
+ name,
244
+ agent_type,
245
+ version,
246
+ capabilities,
247
+ *args,
248
+ **kwargs,
249
+ )
250
+
152
251
  return wrapper
153
252
 
154
253
  return decorator
155
254
 
156
- def _trace_sync_agent_execution(self, func, name, agent_type, version, capabilities, *args, **kwargs):
255
+ def _trace_sync_agent_execution(
256
+ self, func, name, agent_type, version, capabilities, *args, **kwargs
257
+ ):
157
258
  # Generate a unique hash_id for this execution context
158
259
  hash_id = str(uuid.uuid4())
159
260
 
@@ -161,6 +262,9 @@ class AgentTracerMixin:
161
262
  if not self.is_active:
162
263
  return func(*args, **kwargs)
163
264
 
265
+ if not self.auto_instrument_agent:
266
+ return func(*args, **kwargs)
267
+
164
268
  start_time = datetime.now()
165
269
  self.start_time = start_time
166
270
  self.input_data = self._sanitize_input(args, kwargs)
@@ -168,11 +272,11 @@ class AgentTracerMixin:
168
272
  component_id = str(uuid.uuid4())
169
273
 
170
274
  # Extract ground truth if present
171
- ground_truth = kwargs.pop('gt', None) if kwargs else None
275
+ ground_truth = kwargs.pop("gt", None) if kwargs else None
172
276
 
173
277
  # Get parent agent ID if exists
174
278
  parent_agent_id = self.current_agent_id.get()
175
-
279
+
176
280
  # Set the current agent context
177
281
  agent_token = self.current_agent_id.set(component_id)
178
282
  agent_name_token = self.current_agent_name.set(name)
@@ -180,7 +284,7 @@ class AgentTracerMixin:
180
284
  # Initialize empty children list for this agent
181
285
  parent_children = self.agent_children.get()
182
286
  children_token = self.agent_children.set([])
183
-
287
+
184
288
  # Start tracking network calls for this component
185
289
  self.start_component(component_id)
186
290
 
@@ -211,12 +315,12 @@ class AgentTracerMixin:
211
315
  input_data=self.input_data,
212
316
  output_data=self._sanitize_output(result),
213
317
  children=children,
214
- parent_id=parent_agent_id
318
+ parent_id=parent_agent_id,
215
319
  )
216
320
  # Add ground truth to component data if present
217
321
  if ground_truth is not None:
218
322
  agent_component["data"]["gt"] = ground_truth
219
-
323
+
220
324
  # Add this component as a child to parent's children list
221
325
  parent_children.append(agent_component)
222
326
  self.agent_children.set(parent_children)
@@ -231,19 +335,19 @@ class AgentTracerMixin:
231
335
  "code": 500,
232
336
  "type": type(e).__name__,
233
337
  "message": str(e),
234
- "details": {}
338
+ "details": {},
235
339
  }
236
-
340
+
237
341
  # Get children even in case of error
238
342
  children = self.agent_children.get()
239
-
343
+
240
344
  # Set parent_id for all children
241
345
  for child in children:
242
346
  child["parent_id"] = component_id
243
-
347
+
244
348
  # End tracking network calls for this component
245
349
  self.end_component(component_id)
246
-
350
+
247
351
  agent_component = self.create_agent_component(
248
352
  component_id=component_id,
249
353
  hash_id=hash_id,
@@ -257,39 +361,45 @@ class AgentTracerMixin:
257
361
  output_data=None,
258
362
  error=error_component,
259
363
  children=children,
260
- parent_id=parent_agent_id # Add parent ID if exists
364
+ parent_id=parent_agent_id, # Add parent ID if exists
261
365
  )
262
366
  # If this is a nested agent, add it to parent's children
263
367
  if parent_agent_id:
264
368
  parent_component = self._agent_components.get(parent_agent_id)
265
369
  if parent_component:
266
- if 'children' not in parent_component['data']:
267
- parent_component['data']['children'] = []
268
- parent_component['data']['children'].append(agent_component)
370
+ if "children" not in parent_component["data"]:
371
+ parent_component["data"]["children"] = []
372
+ parent_component["data"]["children"].append(agent_component)
269
373
  else:
270
374
  # Only add to root components if no parent
271
375
  self.add_component(agent_component)
376
+
272
377
  raise
273
378
  finally:
274
379
  self.current_agent_id.reset(agent_token)
275
380
  self.current_agent_name.reset(agent_name_token)
276
381
  self.agent_children.reset(children_token)
277
382
 
278
- async def _trace_agent_execution(self, func, name, agent_type, version, capabilities, hash_id, *args, **kwargs):
383
+ async def _trace_agent_execution(
384
+ self, func, name, agent_type, version, capabilities, hash_id, *args, **kwargs
385
+ ):
279
386
  """Asynchronous version of agent tracing"""
280
387
  if not self.is_active:
281
388
  return await func(*args, **kwargs)
282
389
 
390
+ if not self.auto_instrument_agent:
391
+ return await func(*args, **kwargs)
392
+
283
393
  start_time = datetime.now()
284
394
  start_memory = psutil.Process().memory_info().rss
285
395
  component_id = str(uuid.uuid4())
286
396
 
287
397
  # Extract ground truth if present
288
- ground_truth = kwargs.pop('gt', None) if kwargs else None
398
+ ground_truth = kwargs.pop("gt", None) if kwargs else None
289
399
 
290
400
  # Get parent agent ID if exists
291
401
  parent_agent_id = self.current_agent_id.get()
292
-
402
+
293
403
  # Set the current agent context
294
404
  agent_token = self.current_agent_id.set(component_id)
295
405
  agent_name_token = self.current_agent_name.set(name)
@@ -297,6 +407,7 @@ class AgentTracerMixin:
297
407
  # Initialize empty children list for this agent
298
408
  parent_children = self.agent_children.get()
299
409
  children_token = self.agent_children.set([])
410
+ self.start_component(component_id)
300
411
 
301
412
  try:
302
413
  # Execute the agent
@@ -309,6 +420,8 @@ class AgentTracerMixin:
309
420
  # Get children components collected during execution
310
421
  children = self.agent_children.get()
311
422
 
423
+ self.end_component(component_id)
424
+
312
425
  # Create agent component with children and parent if exists
313
426
  agent_component = self.create_agent_component(
314
427
  component_id=component_id,
@@ -322,7 +435,7 @@ class AgentTracerMixin:
322
435
  input_data=self._sanitize_input(args, kwargs),
323
436
  output_data=self._sanitize_output(result),
324
437
  children=children,
325
- parent_id=parent_agent_id
438
+ parent_id=parent_agent_id,
326
439
  )
327
440
 
328
441
  # Add ground truth to component data if present
@@ -343,19 +456,19 @@ class AgentTracerMixin:
343
456
  "code": 500,
344
457
  "type": type(e).__name__,
345
458
  "message": str(e),
346
- "details": {}
459
+ "details": {},
347
460
  }
348
-
461
+
349
462
  # Get children even in case of error
350
463
  children = self.agent_children.get()
351
-
464
+
352
465
  # Set parent_id for all children
353
466
  for child in children:
354
467
  child["parent_id"] = component_id
355
-
468
+
356
469
  # End tracking network calls for this component
357
470
  self.end_component(component_id)
358
-
471
+
359
472
  agent_component = self.create_agent_component(
360
473
  component_id=component_id,
361
474
  hash_id=hash_id,
@@ -369,19 +482,20 @@ class AgentTracerMixin:
369
482
  output_data=None,
370
483
  error=error_component,
371
484
  children=children,
372
- parent_id=parent_agent_id # Add parent ID if exists
485
+ parent_id=parent_agent_id, # Add parent ID if exists
373
486
  )
374
-
487
+
375
488
  # If this is a nested agent, add it to parent's children
376
489
  if parent_agent_id:
377
490
  parent_component = self._agent_components.get(parent_agent_id)
378
491
  if parent_component:
379
- if 'children' not in parent_component['data']:
380
- parent_component['data']['children'] = []
381
- parent_component['data']['children'].append(agent_component)
492
+ if "children" not in parent_component["data"]:
493
+ parent_component["data"]["children"] = []
494
+ parent_component["data"]["children"].append(agent_component)
382
495
  else:
383
496
  # Only add to root components if no parent
384
497
  self.add_component(agent_component)
498
+
385
499
  raise
386
500
  finally:
387
501
  # Reset context variables
@@ -391,6 +505,14 @@ class AgentTracerMixin:
391
505
 
392
506
  def create_agent_component(self, **kwargs):
393
507
  """Create an agent component according to the data structure"""
508
+ network_calls = []
509
+ if self.auto_instrument_network:
510
+ network_calls = self.component_network_calls.get(kwargs["component_id"], [])
511
+ interactions = []
512
+ if self.auto_instrument_user_interaction:
513
+ interactions = self.component_user_interaction.get(
514
+ kwargs["component_id"], []
515
+ )
394
516
  start_time = kwargs["start_time"]
395
517
  component = {
396
518
  "id": kwargs["component_id"],
@@ -406,20 +528,25 @@ class AgentTracerMixin:
406
528
  "agent_type": kwargs["agent_type"],
407
529
  "version": kwargs["version"],
408
530
  "capabilities": kwargs["capabilities"],
409
- "memory_used": kwargs["memory_used"]
531
+ "memory_used": kwargs["memory_used"],
532
+ "tags": self.span_attributes_dict[kwargs["name"]].tags or [],
410
533
  },
411
534
  "data": {
412
535
  "input": kwargs["input_data"],
413
536
  "output": kwargs["output_data"],
414
- "children": kwargs.get("children", [])
537
+ "children": kwargs.get("children", []),
415
538
  },
416
- "network_calls": self.component_network_calls.get(kwargs["component_id"], []),
417
- "interactions": self.component_user_interaction.get(kwargs["component_id"], [])
539
+ "metrics": self.span_attributes_dict[kwargs["name"]].metrics or [],
540
+ "network_calls": network_calls,
541
+ "interactions": interactions,
418
542
  }
419
543
 
420
- if self.gt:
544
+ if self.gt:
421
545
  component["data"]["gt"] = self.gt
422
546
 
547
+ # Reset the SpanAttributes context variable
548
+ self.span_attributes_dict[kwargs["name"]] = SpanAttributes(kwargs["name"])
549
+
423
550
  return component
424
551
 
425
552
  def start_component(self, component_id):
@@ -438,21 +565,22 @@ class AgentTracerMixin:
438
565
 
439
566
  def _sanitize_input(self, args: tuple, kwargs: dict) -> str:
440
567
  """Convert input arguments to text format.
441
-
568
+
442
569
  Args:
443
570
  args: Input arguments tuple
444
571
  kwargs: Input keyword arguments dict
445
-
572
+
446
573
  Returns:
447
574
  str: Text representation of the input arguments
448
575
  """
576
+
449
577
  def _sanitize_value(value):
450
578
  if isinstance(value, dict):
451
579
  return str({k: _sanitize_value(v) for k, v in value.items()})
452
580
  elif isinstance(value, (list, tuple)):
453
581
  return str([_sanitize_value(item) for item in value])
454
582
  return str(value)
455
-
583
+
456
584
  sanitized_args = [_sanitize_value(arg) for arg in args]
457
585
  sanitized_kwargs = {k: _sanitize_value(v) for k, v in kwargs.items()}
458
586
  return str({"args": sanitized_args, "kwargs": sanitized_kwargs})
@@ -461,4 +589,13 @@ class AgentTracerMixin:
461
589
  """Sanitize and format output data"""
462
590
  if isinstance(output, (int, float, bool, str, list, dict)):
463
591
  return output
464
- return str(output)
592
+ return str(output)
593
+
594
+ def instrument_agent_calls(self):
595
+ self.auto_instrument_agent = True
596
+
597
+ def instrument_user_interaction_calls(self):
598
+ self.auto_instrument_user_interaction = True
599
+
600
+ def instrument_network_calls(self):
601
+ self.auto_instrument_network = True