ragaai-catalyst 2.1.3b0__py3-none-any.whl → 2.1.4__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 (27) hide show
  1. ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py +37 -11
  2. ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py +240 -81
  3. ragaai_catalyst/tracers/agentic_tracing/tracers/base.py +632 -114
  4. ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py +316 -0
  5. ragaai_catalyst/tracers/agentic_tracing/tracers/langgraph_tracer.py +0 -0
  6. ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py +229 -82
  7. ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py +214 -59
  8. ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py +16 -14
  9. ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py +147 -28
  10. ragaai_catalyst/tracers/agentic_tracing/tracers/user_interaction_tracer.py +88 -2
  11. ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py +9 -51
  12. ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py +83 -0
  13. ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py +26 -0
  14. ragaai_catalyst/tracers/agentic_tracing/utils/get_user_trace_metrics.py +28 -0
  15. ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py +45 -15
  16. ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json +2520 -2152
  17. ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py +59 -0
  18. ragaai_catalyst/tracers/agentic_tracing/utils/trace_utils.py +23 -0
  19. ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py +284 -15
  20. ragaai_catalyst/tracers/llamaindex_callback.py +5 -5
  21. ragaai_catalyst/tracers/tracer.py +83 -10
  22. ragaai_catalyst/tracers/upload_traces.py +1 -1
  23. ragaai_catalyst-2.1.4.dist-info/METADATA +431 -0
  24. {ragaai_catalyst-2.1.3b0.dist-info → ragaai_catalyst-2.1.4.dist-info}/RECORD +26 -20
  25. ragaai_catalyst-2.1.3b0.dist-info/METADATA +0 -43
  26. {ragaai_catalyst-2.1.3b0.dist-info → ragaai_catalyst-2.1.4.dist-info}/WHEEL +0 -0
  27. {ragaai_catalyst-2.1.3b0.dist-info → ragaai_catalyst-2.1.4.dist-info}/top_level.txt +0 -0
@@ -13,46 +13,113 @@ from .tool_tracer import ToolTracerMixin
13
13
  from .agent_tracer import AgentTracerMixin
14
14
  from .network_tracer import NetworkTracer
15
15
  from .user_interaction_tracer import UserInteractionTracer
16
+ from .custom_tracer import CustomTracerMixin
17
+ from ..utils.span_attributes import SpanAttributes
16
18
 
17
19
  from ..data.data_structure import (
18
- Trace, Metadata, SystemInfo, OSInfo, EnvironmentInfo,
19
- Resources, CPUResource, MemoryResource, DiskResource, NetworkResource,
20
- ResourceInfo, MemoryInfo, DiskInfo, NetworkInfo,
21
- Component, LLMComponent, AgentComponent, ToolComponent,
22
- NetworkCall, Interaction, Error
20
+ Trace,
21
+ Metadata,
22
+ SystemInfo,
23
+ OSInfo,
24
+ EnvironmentInfo,
25
+ Resources,
26
+ CPUResource,
27
+ MemoryResource,
28
+ DiskResource,
29
+ NetworkResource,
30
+ ResourceInfo,
31
+ MemoryInfo,
32
+ DiskInfo,
33
+ NetworkInfo,
34
+ Component,
35
+ LLMComponent,
36
+ AgentComponent,
37
+ ToolComponent,
38
+ NetworkCall,
39
+ Interaction,
40
+ Error,
23
41
  )
24
42
 
25
43
  from ....ragaai_catalyst import RagaAICatalyst
26
44
  from ragaai_catalyst.tracers.upload_traces import UploadTraces
27
45
 
28
- class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMixin):
29
- def __init__(self, user_detail, auto_instrument_llm: bool = True):
46
+
47
+ class AgenticTracing(
48
+ BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMixin, CustomTracerMixin
49
+ ):
50
+ def __init__(self, user_detail, auto_instrumentation=None):
30
51
  # Initialize all parent classes
31
52
  self.user_interaction_tracer = UserInteractionTracer()
32
53
  LLMTracerMixin.__init__(self)
33
54
  ToolTracerMixin.__init__(self)
34
55
  AgentTracerMixin.__init__(self)
35
-
56
+ CustomTracerMixin.__init__(self)
57
+
36
58
  self.project_name = user_detail["project_name"]
37
59
  self.project_id = user_detail["project_id"]
38
- self.dataset_name = user_detail["dataset_name"]
60
+ # self.dataset_name = user_detail["dataset_name"]
39
61
  self.trace_user_detail = user_detail["trace_user_detail"]
40
62
  self.base_url = f"{RagaAICatalyst.BASE_URL}"
41
63
  self.timeout = 10
42
64
 
43
65
  BaseTracer.__init__(self, user_detail)
44
-
45
- self.auto_instrument_llm = auto_instrument_llm
66
+
46
67
  self.tools: Dict[str, Tool] = {}
47
68
  self.call_depth = contextvars.ContextVar("call_depth", default=0)
48
- self.current_component_id = contextvars.ContextVar("current_component_id", default=None)
69
+ self.current_component_id = contextvars.ContextVar(
70
+ "current_component_id", default=None
71
+ )
49
72
  self.network_tracer = NetworkTracer()
50
- self.is_active = False
73
+
74
+ # Handle auto_instrumentation
75
+ if auto_instrumentation is None:
76
+ # Default behavior: everything enabled
77
+ self.is_active = True
78
+ self.auto_instrument_llm = True
79
+ self.auto_instrument_tool = True
80
+ self.auto_instrument_agent = True
81
+ self.auto_instrument_user_interaction = True
82
+ self.auto_instrument_file_io = True
83
+ self.auto_instrument_network = True
84
+ self.auto_instrument_custom = True
85
+ else:
86
+ # Set global active state
87
+ self.is_active = (
88
+ any(auto_instrumentation.values())
89
+ if isinstance(auto_instrumentation, dict)
90
+ else bool(auto_instrumentation)
91
+ )
92
+
93
+ # Set individual components
94
+ if isinstance(auto_instrumentation, dict):
95
+ self.auto_instrument_llm = auto_instrumentation.get("llm", False)
96
+ self.auto_instrument_tool = auto_instrumentation.get("tool", False)
97
+ self.auto_instrument_agent = auto_instrumentation.get("agent", False)
98
+ self.auto_instrument_user_interaction = auto_instrumentation.get(
99
+ "user_interaction", False
100
+ )
101
+ self.auto_instrument_file_io = auto_instrumentation.get(
102
+ "file_io", False
103
+ )
104
+ self.auto_instrument_network = auto_instrumentation.get(
105
+ "network", False
106
+ )
107
+ self.auto_instrument_custom = auto_instrumentation.get("custom", False)
108
+ else:
109
+ # If boolean provided, apply to all components
110
+ self.auto_instrument_llm = bool(auto_instrumentation)
111
+ self.auto_instrument_tool = bool(auto_instrumentation)
112
+ self.auto_instrument_agent = bool(auto_instrumentation)
113
+ self.auto_instrument_user_interaction = bool(auto_instrumentation)
114
+ self.auto_instrument_file_io = bool(auto_instrumentation)
115
+ self.auto_instrument_network = bool(auto_instrumentation)
116
+ self.auto_instrument_custom = bool(auto_instrumentation)
117
+
51
118
  self.current_agent_id = contextvars.ContextVar("current_agent_id", default=None)
52
119
  self.agent_children = contextvars.ContextVar("agent_children", default=[])
53
120
  self.component_network_calls = {} # Store network calls per component
54
121
  self.component_user_interaction = {}
55
-
122
+
56
123
  # Create output directory if it doesn't exist
57
124
  self.output_dir = Path("./traces") # Using default traces directory
58
125
  self.output_dir.mkdir(exist_ok=True)
@@ -61,19 +128,28 @@ class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMix
61
128
  """Start tracking network calls for a component"""
62
129
  self.component_network_calls[component_id] = []
63
130
  self.network_tracer.network_calls = [] # Reset network calls
64
- self.component_user_interaction[component_id] = []
65
131
  self.current_component_id.set(component_id)
66
132
  self.user_interaction_tracer.component_id.set(component_id)
67
133
 
68
134
  def end_component(self, component_id: str):
69
135
  """End tracking network calls for a component"""
70
- self.component_network_calls[component_id] = self.network_tracer.network_calls.copy()
136
+ self.component_network_calls[component_id] = (
137
+ self.network_tracer.network_calls.copy()
138
+ )
71
139
  self.network_tracer.network_calls = [] # Reset for next component
72
- self.component_user_interaction[component_id] = self.user_interaction_tracer.interactions.copy()
73
- self.user_interaction_tracer.interactions = []
140
+ # self.component_user_interaction[component_id] = [interaction for interaction in self.user_interaction_tracer.interactions if interaction.get('component_id') == component_id]
141
+ for interaction in self.user_interaction_tracer.interactions:
142
+ interaction_component_id = interaction.get("component_id")
143
+ if interaction_component_id not in self.component_user_interaction:
144
+ self.component_user_interaction[interaction_component_id] = []
145
+ if interaction not in self.component_user_interaction[interaction_component_id]:
146
+ self.component_user_interaction[interaction_component_id].append(interaction)
74
147
 
75
148
  def start(self):
76
149
  """Start tracing"""
150
+ if not self.is_active:
151
+ return
152
+
77
153
  # Setup user interaction tracing
78
154
  self.user_interaction_tracer.project_id.set(self.project_id)
79
155
  self.user_interaction_tracer.trace_id.set(self.trace_id)
@@ -81,84 +157,115 @@ class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMix
81
157
  self.user_interaction_tracer.component_id.set(self.current_component_id.get())
82
158
  builtins.print = self.user_interaction_tracer.traced_print
83
159
  builtins.input = self.user_interaction_tracer.traced_input
84
-
160
+ builtins.open = self.user_interaction_tracer.traced_open
161
+
85
162
  # Start base tracer (includes system info and resource monitoring)
86
163
  super().start()
87
- self.is_active = True
88
-
164
+
89
165
  # Activate network tracing
90
166
  self.network_tracer.activate_patches()
91
-
92
- # Instrument calls from mixins
167
+
168
+ # take care of the auto instrumentation
93
169
  if self.auto_instrument_llm:
94
170
  self.instrument_llm_calls()
95
171
 
172
+ if self.auto_instrument_tool:
173
+ self.instrument_tool_calls()
174
+
175
+ if self.auto_instrument_agent:
176
+ self.instrument_agent_calls()
177
+
178
+ if self.auto_instrument_custom:
179
+ self.instrument_custom_calls()
180
+
181
+ if self.auto_instrument_user_interaction:
182
+
183
+ ToolTracerMixin.instrument_user_interaction_calls(self)
184
+ LLMTracerMixin.instrument_user_interaction_calls(self)
185
+ AgentTracerMixin.instrument_user_interaction_calls(self)
186
+ CustomTracerMixin.instrument_user_interaction_calls(self)
187
+
188
+ if self.auto_instrument_network:
189
+ ToolTracerMixin.instrument_network_calls(self)
190
+ LLMTracerMixin.instrument_network_calls(self)
191
+ AgentTracerMixin.instrument_network_calls(self)
192
+ CustomTracerMixin.instrument_network_calls(self)
193
+
194
+ # These will be implemented later
195
+ # if self.auto_instrument_file_io:
196
+ # self.instrument_file_io_calls()
197
+
96
198
  def stop(self):
97
199
  """Stop tracing and save results"""
98
200
  if self.is_active:
99
201
  # Restore original print and input functions
100
202
  builtins.print = self.user_interaction_tracer.original_print
101
203
  builtins.input = self.user_interaction_tracer.original_input
102
-
204
+ builtins.open = self.user_interaction_tracer.original_open
205
+
103
206
  # Calculate final metrics before stopping
104
207
  self._calculate_final_metrics()
105
-
208
+
106
209
  # Deactivate network tracing
107
210
  self.network_tracer.deactivate_patches()
108
-
211
+
212
+ # Clear visited metrics when stopping trace
213
+ self.visited_metrics.clear()
214
+
109
215
  # Stop base tracer (includes saving to file)
110
216
  super().stop()
111
-
217
+
112
218
  # Cleanup
113
219
  self.unpatch_llm_calls()
114
220
  self.user_interaction_tracer.interactions = [] # Clear interactions list
115
221
  self.is_active = False
116
222
 
117
-
118
223
  def _calculate_final_metrics(self):
119
224
  """Calculate total cost and tokens from all components"""
120
225
  total_cost = 0.0
121
226
  total_tokens = 0
122
-
227
+
123
228
  def process_component(component):
124
229
  nonlocal total_cost, total_tokens
125
230
  # Convert component to dict if it's an object
126
- comp_dict = component.__dict__ if hasattr(component, '__dict__') else component
127
-
128
- if comp_dict.get('type') == "llm":
129
- info = comp_dict.get('info', {})
231
+ comp_dict = (
232
+ component.__dict__ if hasattr(component, "__dict__") else component
233
+ )
234
+
235
+ if comp_dict.get("type") == "llm":
236
+ info = comp_dict.get("info", {})
130
237
  if isinstance(info, dict):
131
238
  # Extract cost
132
- cost_info = info.get('cost', {})
239
+ cost_info = info.get("cost", {})
133
240
  if isinstance(cost_info, dict):
134
- total_cost += cost_info.get('total_cost', 0)
135
-
241
+ total_cost += cost_info.get("total_cost", 0)
242
+
136
243
  # Extract tokens
137
- token_info = info.get('tokens', {})
244
+ token_info = info.get("tokens", {})
138
245
  if isinstance(token_info, dict):
139
- total_tokens += token_info.get('total_tokens', 0)
246
+ total_tokens += token_info.get("total_tokens", 0)
140
247
  else:
141
- token_info = info.get('token_usage', {})
248
+ token_info = info.get("token_usage", {})
142
249
  if isinstance(token_info, dict):
143
- total_tokens += token_info.get('total_tokens', 0)
144
-
250
+ total_tokens += token_info.get("total_tokens", 0)
251
+
145
252
  # Process children if they exist
146
- data = comp_dict.get('data', {})
253
+ data = comp_dict.get("data", {})
147
254
  if isinstance(data, dict):
148
- children = data.get('children', [])
255
+ children = data.get("children", [])
149
256
  if children:
150
257
  for child in children:
151
258
  process_component(child)
152
-
259
+
153
260
  # Process all root components
154
261
  for component in self.components:
155
262
  process_component(component)
156
-
263
+
157
264
  # Update metadata in trace
158
- if hasattr(self, 'trace'):
265
+ if hasattr(self, "trace"):
159
266
  if isinstance(self.trace.metadata, dict):
160
- self.trace.metadata['total_cost'] = total_cost
161
- self.trace.metadata['total_tokens'] = total_tokens
267
+ self.trace.metadata["total_cost"] = total_cost
268
+ self.trace.metadata["total_tokens"] = total_tokens
162
269
  else:
163
270
  self.trace.metadata.total_cost = total_cost
164
271
  self.trace.metadata.total_tokens = total_tokens
@@ -166,11 +273,34 @@ class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMix
166
273
  def add_component(self, component_data: dict, is_error: bool = False):
167
274
  """Add a component to the trace data"""
168
275
  # Convert dict to appropriate Component type
169
- filtered_data = {k: v for k, v in component_data.items() if k in ["id", "hash_id", "source_hash_id", "type", "name", "start_time", "end_time", "parent_id", "info", "extra_info", "data", "network_calls", "interactions", "error"]}
276
+ filtered_data = {
277
+ k: v
278
+ for k, v in component_data.items()
279
+ if k
280
+ in [
281
+ "id",
282
+ "hash_id",
283
+ "source_hash_id",
284
+ "type",
285
+ "name",
286
+ "start_time",
287
+ "end_time",
288
+ "parent_id",
289
+ "info",
290
+ "extra_info",
291
+ "data",
292
+ "metadata",
293
+ "metrics",
294
+ "feedback",
295
+ "network_calls",
296
+ "interactions",
297
+ "error",
298
+ ]
299
+ }
170
300
 
171
301
  if component_data["type"] == "llm":
172
302
  component = LLMComponent(**filtered_data)
173
- elif component_data["type"] == "agent":
303
+ elif component_data["type"] == "agent":
174
304
  component = AgentComponent(**filtered_data)
175
305
  elif component_data["type"] == "tool":
176
306
  component = ToolComponent(**filtered_data)
@@ -187,21 +317,27 @@ class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMix
187
317
  else:
188
318
  # Add component to the main trace
189
319
  super().add_component(component)
190
-
320
+
191
321
  # Handle error case
192
322
  if is_error:
193
323
  # Get the parent component if it exists
194
324
  parent_id = component_data.get("parent_id")
195
325
  children = self.agent_children.get()
196
-
326
+
197
327
  # Set parent_id for all children
198
328
  for child in children:
199
329
  child["parent_id"] = parent_id
200
-
330
+
201
331
  agent_tracer_mixin = AgentTracerMixin()
202
332
  agent_tracer_mixin.component_network_calls = self.component_network_calls
203
- agent_tracer_mixin.component_user_interaction = self.component_user_interaction
204
-
333
+ agent_tracer_mixin.component_user_interaction = (
334
+ self.component_user_interaction
335
+ )
336
+
337
+ agent_tracer_mixin.span_attributes_dict[self.current_agent_name.get()] = (
338
+ SpanAttributes(self.current_agent_name.get())
339
+ )
340
+
205
341
  # Create parent component with error info
206
342
  parent_component = agent_tracer_mixin.create_agent_component(
207
343
  component_id=parent_id,
@@ -213,15 +349,34 @@ class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMix
213
349
  version=self.version.get(),
214
350
  capabilities=self.capabilities.get(),
215
351
  start_time=self.start_time,
216
- end_time=datetime.now(),
352
+ end_time=datetime.now().astimezone().isoformat(),
217
353
  memory_used=0,
218
354
  input_data=self.input_data,
219
355
  output_data=None,
220
356
  children=children,
221
- parent_id=None # Add parent ID if exists
357
+ parent_id=None, # Add parent ID if exists
222
358
  )
223
359
 
224
- filtered_data = {k: v for k, v in parent_component.items() if k in ["id", "hash_id", "source_hash_id", "type", "name", "start_time", "end_time", "parent_id", "info", "data", "network_calls", "interactions", "error"]}
360
+ filtered_data = {
361
+ k: v
362
+ for k, v in parent_component.items()
363
+ if k
364
+ in [
365
+ "id",
366
+ "hash_id",
367
+ "source_hash_id",
368
+ "type",
369
+ "name",
370
+ "start_time",
371
+ "end_time",
372
+ "parent_id",
373
+ "info",
374
+ "data",
375
+ "network_calls",
376
+ "interactions",
377
+ "error",
378
+ ]
379
+ }
225
380
  parent_agent_component = AgentComponent(**filtered_data)
226
381
  # Add the parent component to trace and stop tracing
227
382
  super().add_component(parent_agent_component)
@@ -234,4 +389,4 @@ class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMix
234
389
 
235
390
  def __exit__(self, exc_type, exc_value, traceback):
236
391
  """Context manager exit"""
237
- self.stop()
392
+ self.stop()
@@ -47,6 +47,8 @@ class NetworkTracer:
47
47
  "url": url,
48
48
  "method": method,
49
49
  "status_code": status_code,
50
+ "start_time": start_time.isoformat() if start_time else None,
51
+ "end_time": end_time.isoformat() if end_time else None,
50
52
  "response_time": duration,
51
53
  "bytes_sent": bytes_sent,
52
54
  "bytes_received": bytes_received,
@@ -103,10 +105,10 @@ def monkey_patch_urllib(network_tracer):
103
105
  method = url.get_method()
104
106
  url_str = url.full_url
105
107
 
106
- start_time = datetime.now()
108
+ start_time = datetime.now().astimezone()
107
109
  try:
108
110
  response = original_urlopen(url, data, timeout, *args, **kwargs)
109
- end_time = datetime.now()
111
+ end_time = datetime.now().astimezone()
110
112
  network_tracer.record_call(
111
113
  method=method,
112
114
  url=url_str,
@@ -120,7 +122,7 @@ def monkey_patch_urllib(network_tracer):
120
122
  )
121
123
  return response
122
124
  except Exception as e:
123
- end_time = datetime.now()
125
+ end_time = datetime.now().astimezone()
124
126
  network_tracer.record_call(
125
127
  method=method,
126
128
  url=url_str,
@@ -142,10 +144,10 @@ def monkey_patch_requests(network_tracer):
142
144
  original_request = requests.Session.request
143
145
 
144
146
  def patched_request(self, method, url, *args, **kwargs):
145
- start_time = datetime.now()
147
+ start_time = datetime.now().astimezone()
146
148
  try:
147
149
  response = original_request(self, method, url, *args, **kwargs)
148
- end_time = datetime.now()
150
+ end_time = datetime.now().astimezone()
149
151
  network_tracer.record_call(
150
152
  method=method,
151
153
  url=url,
@@ -159,7 +161,7 @@ def monkey_patch_requests(network_tracer):
159
161
  )
160
162
  return response
161
163
  except Exception as e:
162
- end_time = datetime.now()
164
+ end_time = datetime.now().astimezone()
163
165
  network_tracer.record_call(
164
166
  method=method,
165
167
  url=url,
@@ -182,7 +184,7 @@ def monkey_patch_http_client(network_tracer):
182
184
  original_https_request = HTTPSConnection.request
183
185
 
184
186
  def patched_request(self, method, url, body=None, headers=None, *args, **kwargs):
185
- start_time = datetime.now()
187
+ start_time = datetime.now().astimezone()
186
188
  try:
187
189
  result = (
188
190
  original_http_request(self, method, url, body, headers, *args, **kwargs)
@@ -192,7 +194,7 @@ def monkey_patch_http_client(network_tracer):
192
194
  )
193
195
  )
194
196
  response = self.getresponse()
195
- end_time = datetime.now()
197
+ end_time = datetime.now().astimezone()
196
198
  network_tracer.record_call(
197
199
  method=method,
198
200
  url=f"{self._http_vsn_str} {self.host}:{self.port}{url}",
@@ -206,7 +208,7 @@ def monkey_patch_http_client(network_tracer):
206
208
  )
207
209
  return result
208
210
  except Exception as e:
209
- end_time = datetime.now()
211
+ end_time = datetime.now().astimezone()
210
212
  network_tracer.record_call(
211
213
  method=method,
212
214
  url=f"{self._http_vsn_str} {self.host}:{self.port}{url}",
@@ -231,10 +233,10 @@ def monkey_patch_socket(network_tracer):
231
233
 
232
234
  def patched_create_connection(address, *args, **kwargs):
233
235
  host, port = address
234
- start_time = datetime.now()
236
+ start_time = datetime.now().astimezone()
235
237
  try:
236
238
  result = original_create_connection(address, *args, **kwargs)
237
- end_time = datetime.now()
239
+ end_time = datetime.now().astimezone()
238
240
  network_tracer.record_call(
239
241
  method="CONNECT",
240
242
  url=f"{host}:{port}",
@@ -243,7 +245,7 @@ def monkey_patch_socket(network_tracer):
243
245
  )
244
246
  return result
245
247
  except Exception as e:
246
- end_time = datetime.now()
248
+ end_time = datetime.now().astimezone()
247
249
  network_tracer.record_call(
248
250
  method="CONNECT",
249
251
  url=f"{host}:{port}",
@@ -263,10 +265,10 @@ def restore_socket(original_create_connection):
263
265
 
264
266
  async def patch_aiohttp_trace_config(network_tracer):
265
267
  async def on_request_start(session, trace_config_ctx, params):
266
- trace_config_ctx.start = datetime.now()
268
+ trace_config_ctx.start = datetime.now().astimezone()
267
269
 
268
270
  async def on_request_end(session, trace_config_ctx, params):
269
- end_time = datetime.now()
271
+ end_time = datetime.now().astimezone()
270
272
  response = params.response
271
273
  network_tracer.record_call(
272
274
  method=params.method,