mseep-agentops 0.4.18__py3-none-any.whl → 0.4.23__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 (153) hide show
  1. agentops/__init__.py +0 -0
  2. agentops/client/api/base.py +28 -30
  3. agentops/client/api/versions/v3.py +29 -25
  4. agentops/client/api/versions/v4.py +87 -46
  5. agentops/client/client.py +98 -29
  6. agentops/client/http/README.md +87 -0
  7. agentops/client/http/http_client.py +126 -172
  8. agentops/config.py +8 -2
  9. agentops/instrumentation/OpenTelemetry.md +133 -0
  10. agentops/instrumentation/README.md +167 -0
  11. agentops/instrumentation/__init__.py +13 -1
  12. agentops/instrumentation/agentic/ag2/__init__.py +18 -0
  13. agentops/instrumentation/agentic/ag2/instrumentor.py +922 -0
  14. agentops/instrumentation/agentic/agno/__init__.py +19 -0
  15. agentops/instrumentation/agentic/agno/attributes/__init__.py +20 -0
  16. agentops/instrumentation/agentic/agno/attributes/agent.py +250 -0
  17. agentops/instrumentation/agentic/agno/attributes/metrics.py +214 -0
  18. agentops/instrumentation/agentic/agno/attributes/storage.py +158 -0
  19. agentops/instrumentation/agentic/agno/attributes/team.py +195 -0
  20. agentops/instrumentation/agentic/agno/attributes/tool.py +210 -0
  21. agentops/instrumentation/agentic/agno/attributes/workflow.py +254 -0
  22. agentops/instrumentation/agentic/agno/instrumentor.py +1313 -0
  23. agentops/instrumentation/agentic/crewai/LICENSE +201 -0
  24. agentops/instrumentation/agentic/crewai/NOTICE.md +10 -0
  25. agentops/instrumentation/agentic/crewai/__init__.py +6 -0
  26. agentops/instrumentation/agentic/crewai/crewai_span_attributes.py +335 -0
  27. agentops/instrumentation/agentic/crewai/instrumentation.py +535 -0
  28. agentops/instrumentation/agentic/crewai/version.py +1 -0
  29. agentops/instrumentation/agentic/google_adk/__init__.py +19 -0
  30. agentops/instrumentation/agentic/google_adk/instrumentor.py +68 -0
  31. agentops/instrumentation/agentic/google_adk/patch.py +767 -0
  32. agentops/instrumentation/agentic/haystack/__init__.py +1 -0
  33. agentops/instrumentation/agentic/haystack/instrumentor.py +186 -0
  34. agentops/instrumentation/agentic/langgraph/__init__.py +3 -0
  35. agentops/instrumentation/agentic/langgraph/attributes.py +54 -0
  36. agentops/instrumentation/agentic/langgraph/instrumentation.py +598 -0
  37. agentops/instrumentation/agentic/langgraph/version.py +1 -0
  38. agentops/instrumentation/agentic/openai_agents/README.md +156 -0
  39. agentops/instrumentation/agentic/openai_agents/SPANS.md +145 -0
  40. agentops/instrumentation/agentic/openai_agents/TRACING_API.md +144 -0
  41. agentops/instrumentation/agentic/openai_agents/__init__.py +30 -0
  42. agentops/instrumentation/agentic/openai_agents/attributes/common.py +549 -0
  43. agentops/instrumentation/agentic/openai_agents/attributes/completion.py +172 -0
  44. agentops/instrumentation/agentic/openai_agents/attributes/model.py +58 -0
  45. agentops/instrumentation/agentic/openai_agents/attributes/tokens.py +275 -0
  46. agentops/instrumentation/agentic/openai_agents/exporter.py +469 -0
  47. agentops/instrumentation/agentic/openai_agents/instrumentor.py +107 -0
  48. agentops/instrumentation/agentic/openai_agents/processor.py +58 -0
  49. agentops/instrumentation/agentic/smolagents/README.md +88 -0
  50. agentops/instrumentation/agentic/smolagents/__init__.py +12 -0
  51. agentops/instrumentation/agentic/smolagents/attributes/agent.py +354 -0
  52. agentops/instrumentation/agentic/smolagents/attributes/model.py +205 -0
  53. agentops/instrumentation/agentic/smolagents/instrumentor.py +286 -0
  54. agentops/instrumentation/agentic/smolagents/stream_wrapper.py +258 -0
  55. agentops/instrumentation/agentic/xpander/__init__.py +15 -0
  56. agentops/instrumentation/agentic/xpander/context.py +112 -0
  57. agentops/instrumentation/agentic/xpander/instrumentor.py +877 -0
  58. agentops/instrumentation/agentic/xpander/trace_probe.py +86 -0
  59. agentops/instrumentation/agentic/xpander/version.py +3 -0
  60. agentops/instrumentation/common/README.md +65 -0
  61. agentops/instrumentation/common/attributes.py +1 -2
  62. agentops/instrumentation/providers/anthropic/__init__.py +24 -0
  63. agentops/instrumentation/providers/anthropic/attributes/__init__.py +23 -0
  64. agentops/instrumentation/providers/anthropic/attributes/common.py +64 -0
  65. agentops/instrumentation/providers/anthropic/attributes/message.py +541 -0
  66. agentops/instrumentation/providers/anthropic/attributes/tools.py +231 -0
  67. agentops/instrumentation/providers/anthropic/event_handler_wrapper.py +90 -0
  68. agentops/instrumentation/providers/anthropic/instrumentor.py +146 -0
  69. agentops/instrumentation/providers/anthropic/stream_wrapper.py +436 -0
  70. agentops/instrumentation/providers/google_genai/README.md +33 -0
  71. agentops/instrumentation/providers/google_genai/__init__.py +24 -0
  72. agentops/instrumentation/providers/google_genai/attributes/__init__.py +25 -0
  73. agentops/instrumentation/providers/google_genai/attributes/chat.py +125 -0
  74. agentops/instrumentation/providers/google_genai/attributes/common.py +88 -0
  75. agentops/instrumentation/providers/google_genai/attributes/model.py +284 -0
  76. agentops/instrumentation/providers/google_genai/instrumentor.py +170 -0
  77. agentops/instrumentation/providers/google_genai/stream_wrapper.py +238 -0
  78. agentops/instrumentation/providers/ibm_watsonx_ai/__init__.py +28 -0
  79. agentops/instrumentation/providers/ibm_watsonx_ai/attributes/__init__.py +27 -0
  80. agentops/instrumentation/providers/ibm_watsonx_ai/attributes/attributes.py +277 -0
  81. agentops/instrumentation/providers/ibm_watsonx_ai/attributes/common.py +104 -0
  82. agentops/instrumentation/providers/ibm_watsonx_ai/instrumentor.py +162 -0
  83. agentops/instrumentation/providers/ibm_watsonx_ai/stream_wrapper.py +302 -0
  84. agentops/instrumentation/providers/mem0/__init__.py +45 -0
  85. agentops/instrumentation/providers/mem0/common.py +377 -0
  86. agentops/instrumentation/providers/mem0/instrumentor.py +270 -0
  87. agentops/instrumentation/providers/mem0/memory.py +430 -0
  88. agentops/instrumentation/providers/openai/__init__.py +21 -0
  89. agentops/instrumentation/providers/openai/attributes/__init__.py +7 -0
  90. agentops/instrumentation/providers/openai/attributes/common.py +55 -0
  91. agentops/instrumentation/providers/openai/attributes/response.py +607 -0
  92. agentops/instrumentation/providers/openai/config.py +36 -0
  93. agentops/instrumentation/providers/openai/instrumentor.py +312 -0
  94. agentops/instrumentation/providers/openai/stream_wrapper.py +941 -0
  95. agentops/instrumentation/providers/openai/utils.py +44 -0
  96. agentops/instrumentation/providers/openai/v0.py +176 -0
  97. agentops/instrumentation/providers/openai/v0_wrappers.py +483 -0
  98. agentops/instrumentation/providers/openai/wrappers/__init__.py +30 -0
  99. agentops/instrumentation/providers/openai/wrappers/assistant.py +277 -0
  100. agentops/instrumentation/providers/openai/wrappers/chat.py +259 -0
  101. agentops/instrumentation/providers/openai/wrappers/completion.py +109 -0
  102. agentops/instrumentation/providers/openai/wrappers/embeddings.py +94 -0
  103. agentops/instrumentation/providers/openai/wrappers/image_gen.py +75 -0
  104. agentops/instrumentation/providers/openai/wrappers/responses.py +191 -0
  105. agentops/instrumentation/providers/openai/wrappers/shared.py +81 -0
  106. agentops/instrumentation/utilities/concurrent_futures/__init__.py +10 -0
  107. agentops/instrumentation/utilities/concurrent_futures/instrumentation.py +206 -0
  108. agentops/integration/callbacks/dspy/__init__.py +11 -0
  109. agentops/integration/callbacks/dspy/callback.py +471 -0
  110. agentops/integration/callbacks/langchain/README.md +59 -0
  111. agentops/integration/callbacks/langchain/__init__.py +15 -0
  112. agentops/integration/callbacks/langchain/callback.py +791 -0
  113. agentops/integration/callbacks/langchain/utils.py +54 -0
  114. agentops/legacy/crewai.md +121 -0
  115. agentops/logging/instrument_logging.py +4 -0
  116. agentops/sdk/README.md +220 -0
  117. agentops/sdk/core.py +75 -32
  118. agentops/sdk/descriptors/classproperty.py +28 -0
  119. agentops/sdk/exporters.py +152 -33
  120. agentops/semconv/README.md +125 -0
  121. agentops/semconv/span_kinds.py +0 -2
  122. agentops/validation.py +102 -63
  123. {mseep_agentops-0.4.18.dist-info → mseep_agentops-0.4.23.dist-info}/METADATA +30 -40
  124. mseep_agentops-0.4.23.dist-info/RECORD +178 -0
  125. {mseep_agentops-0.4.18.dist-info → mseep_agentops-0.4.23.dist-info}/WHEEL +1 -2
  126. mseep_agentops-0.4.18.dist-info/RECORD +0 -94
  127. mseep_agentops-0.4.18.dist-info/top_level.txt +0 -2
  128. tests/conftest.py +0 -10
  129. tests/unit/client/__init__.py +0 -1
  130. tests/unit/client/test_http_adapter.py +0 -221
  131. tests/unit/client/test_http_client.py +0 -206
  132. tests/unit/conftest.py +0 -54
  133. tests/unit/sdk/__init__.py +0 -1
  134. tests/unit/sdk/instrumentation_tester.py +0 -207
  135. tests/unit/sdk/test_attributes.py +0 -392
  136. tests/unit/sdk/test_concurrent_instrumentation.py +0 -468
  137. tests/unit/sdk/test_decorators.py +0 -763
  138. tests/unit/sdk/test_exporters.py +0 -241
  139. tests/unit/sdk/test_factory.py +0 -1188
  140. tests/unit/sdk/test_internal_span_processor.py +0 -397
  141. tests/unit/sdk/test_resource_attributes.py +0 -35
  142. tests/unit/test_config.py +0 -82
  143. tests/unit/test_context_manager.py +0 -777
  144. tests/unit/test_events.py +0 -27
  145. tests/unit/test_host_env.py +0 -54
  146. tests/unit/test_init_py.py +0 -501
  147. tests/unit/test_serialization.py +0 -433
  148. tests/unit/test_session.py +0 -676
  149. tests/unit/test_user_agent.py +0 -34
  150. tests/unit/test_validation.py +0 -405
  151. {tests → agentops/instrumentation/agentic/openai_agents/attributes}/__init__.py +0 -0
  152. /tests/unit/__init__.py → /agentops/instrumentation/providers/openai/attributes/tools.py +0 -0
  153. {mseep_agentops-0.4.18.dist-info → mseep_agentops-0.4.23.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,1313 @@
1
+ """Agno Agent Instrumentation for AgentOps
2
+
3
+ This module provides instrumentation for the Agno Agent library, implementing OpenTelemetry
4
+ instrumentation for agent workflows and LLM model calls.
5
+
6
+ This provides clean visibility into agent workflows and actual tool usage with proper
7
+ parent-child span relationships.
8
+ """
9
+
10
+ from typing import List, Any, Optional, Dict
11
+ from opentelemetry import trace, context as otel_context
12
+ from opentelemetry.trace import Status, StatusCode
13
+ from opentelemetry.metrics import Meter
14
+ import threading
15
+ import json
16
+
17
+ from agentops.logging import logger
18
+ from agentops.instrumentation.common import (
19
+ CommonInstrumentor,
20
+ StandardMetrics,
21
+ InstrumentorConfig,
22
+ )
23
+ from agentops.instrumentation.common.wrappers import WrapConfig
24
+
25
+ # Import attribute handlers
26
+ from agentops.instrumentation.agentic.agno.attributes import (
27
+ get_agent_run_attributes,
28
+ get_metrics_attributes,
29
+ get_team_run_attributes,
30
+ get_tool_execution_attributes,
31
+ get_workflow_run_attributes,
32
+ get_workflow_session_attributes,
33
+ get_storage_read_attributes,
34
+ get_storage_write_attributes,
35
+ )
36
+
37
+
38
+ class StreamingContextManager:
39
+ """Manages span contexts for streaming agent and workflow executions."""
40
+
41
+ def __init__(self):
42
+ self._contexts = {} # context_id -> (span_context, span)
43
+ self._agent_sessions = {} # session_id -> agent_id mapping for context lookup
44
+ self._lock = threading.Lock()
45
+
46
+ def store_context(self, context_id: str, span_context: Any, span: Any) -> None:
47
+ """Store span context for streaming execution."""
48
+ with self._lock:
49
+ self._contexts[context_id] = (span_context, span)
50
+
51
+ def get_context(self, context_id: str) -> Optional[tuple]:
52
+ """Retrieve stored span context."""
53
+ with self._lock:
54
+ return self._contexts.get(context_id)
55
+
56
+ def remove_context(self, context_id: str) -> None:
57
+ """Remove stored context (when streaming completes)."""
58
+ with self._lock:
59
+ self._contexts.pop(context_id, None)
60
+
61
+ def store_agent_session_mapping(self, session_id: str, agent_id: str) -> None:
62
+ """Store mapping between session and agent for context lookup."""
63
+ with self._lock:
64
+ self._agent_sessions[session_id] = agent_id
65
+
66
+ def get_agent_context_by_session(self, session_id: str) -> Optional[tuple]:
67
+ """Get agent context using session ID."""
68
+ with self._lock:
69
+ agent_id = self._agent_sessions.get(session_id)
70
+ if agent_id:
71
+ return self._contexts.get(agent_id)
72
+ return None
73
+
74
+ def clear_all(self) -> None:
75
+ """Clear all stored contexts."""
76
+ with self._lock:
77
+ self._contexts.clear()
78
+ self._agent_sessions.clear()
79
+
80
+
81
+ # Methods to wrap for instrumentation
82
+ # Empty list - all wrapping will be done in _custom_wrap to avoid circular imports
83
+ WRAPPED_METHODS: List[WrapConfig] = []
84
+
85
+
86
+ class StreamingResultWrapper:
87
+ """Wrapper for streaming results that maintains agent span as active throughout iteration."""
88
+
89
+ def __init__(self, original_result, span, agent_id, agent_context, streaming_context_manager):
90
+ self.original_result = original_result
91
+ self.span = span
92
+ self.agent_id = agent_id
93
+ self.agent_context = agent_context
94
+ self.streaming_context_manager = streaming_context_manager
95
+ self._consumed = False
96
+
97
+ def __iter__(self):
98
+ """Return iterator that keeps agent span active during iteration."""
99
+ context_token = otel_context.attach(self.agent_context)
100
+ try:
101
+ # Execute iteration within agent context
102
+ for item in self.original_result:
103
+ # Each item is yielded within the agent span context
104
+ yield item
105
+ finally:
106
+ # Clean up when iteration is complete
107
+ otel_context.detach(context_token)
108
+ if not self._consumed:
109
+ self._consumed = True
110
+ self.span.end()
111
+ self.streaming_context_manager.remove_context(self.agent_id)
112
+
113
+ def __getattr__(self, name):
114
+ """Delegate attribute access to the original result."""
115
+ return getattr(self.original_result, name)
116
+
117
+
118
+ def create_streaming_workflow_wrapper(tracer, streaming_context_manager):
119
+ """Create a streaming-aware wrapper for workflow run methods."""
120
+
121
+ def wrapper(wrapped, instance, args, kwargs):
122
+ # Get workflow ID for context storage
123
+ workflow_id = getattr(instance, "workflow_id", None) or getattr(instance, "id", None) or id(instance)
124
+ workflow_id = str(workflow_id)
125
+
126
+ # Get workflow name for span naming
127
+ workflow_name = getattr(instance, "name", None) or type(instance).__name__
128
+ span_name = f"{workflow_name}.agno.workflow.run.workflow" if workflow_name else "agno.workflow.run.workflow"
129
+
130
+ # Check if streaming is enabled
131
+ is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
132
+
133
+ # For streaming, manually manage span lifecycle
134
+ if is_streaming:
135
+ span = tracer.start_span(span_name)
136
+
137
+ try:
138
+ # Set workflow attributes
139
+ attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
140
+ for key, value in attributes.items():
141
+ span.set_attribute(key, value)
142
+
143
+ # Store context for streaming - capture current context with active span
144
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
145
+ streaming_context_manager.store_context(workflow_id, current_context, span)
146
+
147
+ # Execute the original function within workflow context
148
+ context_token = otel_context.attach(current_context)
149
+ try:
150
+ result = wrapped(*args, **kwargs)
151
+ finally:
152
+ otel_context.detach(context_token)
153
+
154
+ # Set result attributes
155
+ result_attributes = get_workflow_run_attributes(
156
+ args=(instance,) + args, kwargs=kwargs, return_value=result
157
+ )
158
+ for key, value in result_attributes.items():
159
+ if key not in attributes: # Avoid duplicates
160
+ span.set_attribute(key, value)
161
+
162
+ span.set_status(Status(StatusCode.OK))
163
+
164
+ # For streaming results, we need to keep the span open
165
+ # The span will be closed when streaming completes
166
+ return result
167
+
168
+ except Exception as e:
169
+ span.set_status(Status(StatusCode.ERROR, str(e)))
170
+ span.record_exception(e)
171
+ span.end()
172
+ streaming_context_manager.remove_context(workflow_id)
173
+ raise
174
+ else:
175
+ # For non-streaming, use normal context manager
176
+ with tracer.start_as_current_span(span_name) as span:
177
+ try:
178
+ # Set workflow attributes
179
+ attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
180
+ for key, value in attributes.items():
181
+ span.set_attribute(key, value)
182
+
183
+ # Execute the original function
184
+ result = wrapped(*args, **kwargs)
185
+
186
+ # Set result attributes
187
+ result_attributes = get_workflow_run_attributes(
188
+ args=(instance,) + args, kwargs=kwargs, return_value=result
189
+ )
190
+ for key, value in result_attributes.items():
191
+ if key not in attributes: # Avoid duplicates
192
+ span.set_attribute(key, value)
193
+
194
+ span.set_status(Status(StatusCode.OK))
195
+ return result
196
+
197
+ except Exception as e:
198
+ span.set_status(Status(StatusCode.ERROR, str(e)))
199
+ span.record_exception(e)
200
+ raise
201
+
202
+ return wrapper
203
+
204
+
205
+ def create_streaming_workflow_async_wrapper(tracer, streaming_context_manager):
206
+ """Create a streaming-aware async wrapper for workflow run methods."""
207
+
208
+ async def wrapper(wrapped, instance, args, kwargs):
209
+ # Get workflow ID for context storage
210
+ workflow_id = getattr(instance, "workflow_id", None) or getattr(instance, "id", None) or id(instance)
211
+ workflow_id = str(workflow_id)
212
+
213
+ # Get workflow name for span naming
214
+ workflow_name = getattr(instance, "name", None) or type(instance).__name__
215
+ span_name = f"{workflow_name}.agno.workflow.run.workflow" if workflow_name else "agno.workflow.run.workflow"
216
+
217
+ # Check if streaming is enabled
218
+ is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
219
+
220
+ # For streaming, manually manage span lifecycle
221
+ if is_streaming:
222
+ span = tracer.start_span(span_name)
223
+
224
+ try:
225
+ # Set workflow attributes
226
+ attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
227
+ for key, value in attributes.items():
228
+ span.set_attribute(key, value)
229
+
230
+ # Store context for streaming - capture current context with active span
231
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
232
+ streaming_context_manager.store_context(workflow_id, current_context, span)
233
+
234
+ # Execute the original function within workflow context
235
+ context_token = otel_context.attach(current_context)
236
+ try:
237
+ result = await wrapped(*args, **kwargs)
238
+ finally:
239
+ otel_context.detach(context_token)
240
+
241
+ # Set result attributes
242
+ result_attributes = get_workflow_run_attributes(
243
+ args=(instance,) + args, kwargs=kwargs, return_value=result
244
+ )
245
+ for key, value in result_attributes.items():
246
+ if key not in attributes: # Avoid duplicates
247
+ span.set_attribute(key, value)
248
+
249
+ span.set_status(Status(StatusCode.OK))
250
+
251
+ # For streaming results, we need to keep the span open
252
+ # The span will be closed when streaming completes
253
+ return result
254
+
255
+ except Exception as e:
256
+ span.set_status(Status(StatusCode.ERROR, str(e)))
257
+ span.record_exception(e)
258
+ span.end()
259
+ streaming_context_manager.remove_context(workflow_id)
260
+ raise
261
+ else:
262
+ # For non-streaming, use normal context manager
263
+ with tracer.start_as_current_span(span_name) as span:
264
+ try:
265
+ # Set workflow attributes
266
+ attributes = get_workflow_run_attributes(args=(instance,) + args, kwargs=kwargs)
267
+ for key, value in attributes.items():
268
+ span.set_attribute(key, value)
269
+
270
+ # Execute the original function
271
+ result = await wrapped(*args, **kwargs)
272
+
273
+ # Set result attributes
274
+ result_attributes = get_workflow_run_attributes(
275
+ args=(instance,) + args, kwargs=kwargs, return_value=result
276
+ )
277
+ for key, value in result_attributes.items():
278
+ if key not in attributes: # Avoid duplicates
279
+ span.set_attribute(key, value)
280
+
281
+ span.set_status(Status(StatusCode.OK))
282
+ return result
283
+
284
+ except Exception as e:
285
+ span.set_status(Status(StatusCode.ERROR, str(e)))
286
+ span.record_exception(e)
287
+ raise
288
+
289
+ return wrapper
290
+
291
+
292
+ def create_streaming_agent_wrapper(tracer, streaming_context_manager):
293
+ """Create a streaming-aware wrapper for agent run methods."""
294
+
295
+ def wrapper(wrapped, instance, args, kwargs):
296
+ # Get agent ID for context storage
297
+ agent_id = getattr(instance, "agent_id", None) or getattr(instance, "id", None) or id(instance)
298
+ agent_id = str(agent_id)
299
+
300
+ # Get session ID for context mapping
301
+ session_id = getattr(instance, "session_id", None)
302
+
303
+ # Get agent name for span naming
304
+ agent_name = getattr(instance, "name", None)
305
+ span_name = f"{agent_name}.agno.agent.run.agent" if agent_name else "agno.agent.run.agent"
306
+
307
+ # Check if streaming is enabled
308
+ is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
309
+
310
+ # For streaming, manually manage span lifecycle
311
+ if is_streaming:
312
+ span = tracer.start_span(span_name)
313
+
314
+ try:
315
+ # Set agent attributes
316
+ attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
317
+ for key, value in attributes.items():
318
+ span.set_attribute(key, value)
319
+
320
+ # Store context for streaming - capture current context with active span
321
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
322
+ streaming_context_manager.store_context(agent_id, current_context, span)
323
+
324
+ # Store session-to-agent mapping for LLM context lookup
325
+ if session_id:
326
+ streaming_context_manager.store_agent_session_mapping(session_id, agent_id)
327
+
328
+ # Execute the original function within agent context
329
+ context_token = otel_context.attach(current_context)
330
+ try:
331
+ result = wrapped(*args, **kwargs)
332
+ finally:
333
+ otel_context.detach(context_token)
334
+
335
+ # Set result attributes
336
+ result_attributes = get_agent_run_attributes(
337
+ args=(instance,) + args, kwargs=kwargs, return_value=result
338
+ )
339
+ for key, value in result_attributes.items():
340
+ if key not in attributes: # Avoid duplicates
341
+ span.set_attribute(key, value)
342
+
343
+ span.set_status(Status(StatusCode.OK))
344
+
345
+ # Wrap the result to maintain context and end span when complete
346
+ if hasattr(result, "__iter__"):
347
+ return StreamingResultWrapper(result, span, agent_id, current_context, streaming_context_manager)
348
+ else:
349
+ # Not actually streaming, clean up immediately
350
+ span.end()
351
+ streaming_context_manager.remove_context(agent_id)
352
+ return result
353
+
354
+ except Exception as e:
355
+ span.set_status(Status(StatusCode.ERROR, str(e)))
356
+ span.record_exception(e)
357
+ span.end()
358
+ streaming_context_manager.remove_context(agent_id)
359
+ raise
360
+ else:
361
+ # For non-streaming, use normal context manager
362
+ with tracer.start_as_current_span(span_name) as span:
363
+ try:
364
+ # Set agent attributes
365
+ attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
366
+ for key, value in attributes.items():
367
+ span.set_attribute(key, value)
368
+
369
+ # Execute the original function
370
+ result = wrapped(*args, **kwargs)
371
+
372
+ # Set result attributes
373
+ result_attributes = get_agent_run_attributes(
374
+ args=(instance,) + args, kwargs=kwargs, return_value=result
375
+ )
376
+ for key, value in result_attributes.items():
377
+ if key not in attributes: # Avoid duplicates
378
+ span.set_attribute(key, value)
379
+
380
+ span.set_status(Status(StatusCode.OK))
381
+ return result
382
+
383
+ except Exception as e:
384
+ span.set_status(Status(StatusCode.ERROR, str(e)))
385
+ span.record_exception(e)
386
+ raise
387
+
388
+ return wrapper
389
+
390
+
391
+ def create_streaming_agent_async_wrapper(tracer, streaming_context_manager):
392
+ """Create a streaming-aware async wrapper for agent run methods."""
393
+
394
+ async def wrapper(wrapped, instance, args, kwargs):
395
+ # Get agent ID for context storage
396
+ agent_id = getattr(instance, "agent_id", None) or getattr(instance, "id", None) or id(instance)
397
+ agent_id = str(agent_id)
398
+
399
+ # Get session ID for context mapping
400
+ session_id = getattr(instance, "session_id", None)
401
+
402
+ # Get agent name for span naming
403
+ agent_name = getattr(instance, "name", None)
404
+ span_name = f"{agent_name}.agno.agent.run.agent" if agent_name else "agno.agent.run.agent"
405
+
406
+ # Check if streaming is enabled
407
+ is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
408
+
409
+ # For streaming, manually manage span lifecycle
410
+ if is_streaming:
411
+ span = tracer.start_span(span_name)
412
+
413
+ try:
414
+ # Set agent attributes
415
+ attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
416
+ for key, value in attributes.items():
417
+ span.set_attribute(key, value)
418
+
419
+ # Store context for streaming - capture current context with active span
420
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
421
+ streaming_context_manager.store_context(agent_id, current_context, span)
422
+
423
+ # Store session-to-agent mapping for LLM context lookup
424
+ if session_id:
425
+ streaming_context_manager.store_agent_session_mapping(session_id, agent_id)
426
+
427
+ # Execute the original function within agent context
428
+ context_token = otel_context.attach(current_context)
429
+ try:
430
+ result = await wrapped(*args, **kwargs)
431
+ finally:
432
+ otel_context.detach(context_token)
433
+
434
+ # Set result attributes
435
+ result_attributes = get_agent_run_attributes(
436
+ args=(instance,) + args, kwargs=kwargs, return_value=result
437
+ )
438
+ for key, value in result_attributes.items():
439
+ if key not in attributes: # Avoid duplicates
440
+ span.set_attribute(key, value)
441
+
442
+ span.set_status(Status(StatusCode.OK))
443
+
444
+ # Wrap the result to maintain context and end span when complete
445
+ if hasattr(result, "__iter__"):
446
+ return StreamingResultWrapper(result, span, agent_id, current_context, streaming_context_manager)
447
+ else:
448
+ # Not actually streaming, clean up immediately
449
+ span.end()
450
+ streaming_context_manager.remove_context(agent_id)
451
+ return result
452
+
453
+ except Exception as e:
454
+ span.set_status(Status(StatusCode.ERROR, str(e)))
455
+ span.record_exception(e)
456
+ span.end()
457
+ streaming_context_manager.remove_context(agent_id)
458
+ raise
459
+ else:
460
+ # For non-streaming, use normal context manager
461
+ with tracer.start_as_current_span(span_name) as span:
462
+ try:
463
+ # Set agent attributes
464
+ attributes = get_agent_run_attributes(args=(instance,) + args, kwargs=kwargs)
465
+ for key, value in attributes.items():
466
+ span.set_attribute(key, value)
467
+
468
+ # Execute the original function
469
+ result = await wrapped(*args, **kwargs)
470
+
471
+ # Set result attributes
472
+ result_attributes = get_agent_run_attributes(
473
+ args=(instance,) + args, kwargs=kwargs, return_value=result
474
+ )
475
+ for key, value in result_attributes.items():
476
+ if key not in attributes: # Avoid duplicates
477
+ span.set_attribute(key, value)
478
+
479
+ span.set_status(Status(StatusCode.OK))
480
+ return result
481
+
482
+ except Exception as e:
483
+ span.set_status(Status(StatusCode.ERROR, str(e)))
484
+ span.record_exception(e)
485
+ raise
486
+
487
+ return wrapper
488
+
489
+
490
+ def create_streaming_tool_wrapper(tracer, streaming_context_manager):
491
+ """Create a streaming-aware wrapper for tool execution methods."""
492
+
493
+ def wrapper(wrapped, instance, args, kwargs):
494
+ # Try to find the agent or workflow context for proper span hierarchy
495
+ parent_context = None
496
+ parent_span = None
497
+
498
+ # Try to get context from agent
499
+ try:
500
+ if hasattr(instance, "_agent"):
501
+ agent = instance._agent
502
+ agent_id = getattr(agent, "agent_id", None) or getattr(agent, "id", None) or id(agent)
503
+ agent_id = str(agent_id)
504
+ context_info = streaming_context_manager.get_context(agent_id)
505
+ if context_info:
506
+ parent_context, parent_span = context_info
507
+ except Exception:
508
+ pass # Continue without agent context if not found
509
+
510
+ # Try to get context from workflow if agent context not found
511
+ if not parent_context:
512
+ try:
513
+ if hasattr(instance, "_workflow"):
514
+ workflow = instance._workflow
515
+ workflow_id = (
516
+ getattr(workflow, "workflow_id", None) or getattr(workflow, "id", None) or id(workflow)
517
+ )
518
+ workflow_id = str(workflow_id)
519
+ context_info = streaming_context_manager.get_context(workflow_id)
520
+ if context_info:
521
+ parent_context, parent_span = context_info
522
+ except Exception:
523
+ pass # Continue without workflow context if not found
524
+
525
+ # Use parent context if available, otherwise use current context
526
+ if parent_context:
527
+ context_token = otel_context.attach(parent_context)
528
+ try:
529
+ with tracer.start_as_current_span("agno.tool.execute.tool_usage") as span:
530
+ try:
531
+ # Set tool attributes
532
+ attributes = get_tool_execution_attributes(args=(instance,) + args, kwargs=kwargs)
533
+ for key, value in attributes.items():
534
+ span.set_attribute(key, value)
535
+
536
+ # Execute the original function
537
+ result = wrapped(*args, **kwargs)
538
+
539
+ # Set result attributes
540
+ result_attributes = get_tool_execution_attributes(
541
+ args=(instance,) + args, kwargs=kwargs, return_value=result
542
+ )
543
+ for key, value in result_attributes.items():
544
+ if key not in attributes: # Avoid duplicates
545
+ span.set_attribute(key, value)
546
+
547
+ span.set_status(Status(StatusCode.OK))
548
+ return result
549
+
550
+ except Exception as e:
551
+ span.set_status(Status(StatusCode.ERROR, str(e)))
552
+ span.record_exception(e)
553
+ raise
554
+ finally:
555
+ otel_context.detach(context_token)
556
+ else:
557
+ # Fallback to normal span creation
558
+ with tracer.start_as_current_span("agno.tool.execute.tool_usage") as span:
559
+ try:
560
+ # Set tool attributes
561
+ attributes = get_tool_execution_attributes(args=(instance,) + args, kwargs=kwargs)
562
+ for key, value in attributes.items():
563
+ span.set_attribute(key, value)
564
+
565
+ # Execute the original function
566
+ result = wrapped(*args, **kwargs)
567
+
568
+ # Set result attributes
569
+ result_attributes = get_tool_execution_attributes(
570
+ args=(instance,) + args, kwargs=kwargs, return_value=result
571
+ )
572
+ for key, value in result_attributes.items():
573
+ if key not in attributes: # Avoid duplicates
574
+ span.set_attribute(key, value)
575
+
576
+ span.set_status(Status(StatusCode.OK))
577
+ return result
578
+
579
+ except Exception as e:
580
+ span.set_status(Status(StatusCode.ERROR, str(e)))
581
+ span.record_exception(e)
582
+ raise
583
+
584
+ return wrapper
585
+
586
+
587
+ def create_metrics_wrapper(tracer, streaming_context_manager):
588
+ """Create a wrapper for metrics methods with dynamic span naming."""
589
+
590
+ def wrapper(wrapped, instance, args, kwargs):
591
+ # Extract model ID for dynamic span naming
592
+ span_name = "agno.agent.metrics" # fallback
593
+ if hasattr(instance, "model") and instance.model and hasattr(instance.model, "id"):
594
+ model_id = str(instance.model.id)
595
+ span_name = f"{model_id}.llm"
596
+
597
+ with tracer.start_as_current_span(span_name) as span:
598
+ try:
599
+ # Set attributes
600
+ attributes = get_metrics_attributes(args=(instance,) + args, kwargs=kwargs)
601
+ for key, value in attributes.items():
602
+ span.set_attribute(key, value)
603
+
604
+ span.set_status(Status(StatusCode.OK))
605
+
606
+ except Exception as e:
607
+ span.set_status(Status(StatusCode.ERROR, str(e)))
608
+ span.record_exception(e)
609
+ raise
610
+
611
+ return wrapper
612
+
613
+
614
+ def create_team_internal_wrapper(tracer, streaming_context_manager):
615
+ """Create a wrapper for Team internal methods (_run/_arun) that manages team span lifecycle."""
616
+
617
+ def wrapper(wrapped, instance, args, kwargs):
618
+ # Get team ID for context storage
619
+ team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
620
+ team_id = str(team_id)
621
+
622
+ # Get team name for span naming
623
+ team_name = getattr(instance, "name", None)
624
+ span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
625
+
626
+ # Check if we already have a team context (from print_response)
627
+ existing_context = streaming_context_manager.get_context(team_id)
628
+
629
+ if existing_context:
630
+ # We're being called from print_response, use existing context
631
+ parent_context, parent_span = existing_context
632
+
633
+ # Execute within the existing team context
634
+ context_token = otel_context.attach(parent_context)
635
+ try:
636
+ with tracer.start_as_current_span(span_name) as span:
637
+ try:
638
+ # Set workflow attributes
639
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
640
+ for key, value in attributes.items():
641
+ span.set_attribute(key, value)
642
+
643
+ # Execute the original function
644
+ result = wrapped(*args, **kwargs)
645
+
646
+ span.set_status(Status(StatusCode.OK))
647
+ return result
648
+
649
+ except Exception as e:
650
+ span.set_status(Status(StatusCode.ERROR, str(e)))
651
+ span.record_exception(e)
652
+ raise
653
+ finally:
654
+ # Close the parent team span when workflow completes
655
+ if parent_span:
656
+ parent_span.end()
657
+ streaming_context_manager.remove_context(team_id)
658
+ finally:
659
+ otel_context.detach(context_token)
660
+ else:
661
+ # Direct call to _run, create new team span
662
+ with tracer.start_as_current_span(span_name) as span:
663
+ try:
664
+ # Set workflow attributes
665
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
666
+ for key, value in attributes.items():
667
+ span.set_attribute(key, value)
668
+
669
+ # Execute the original function
670
+ result = wrapped(*args, **kwargs)
671
+
672
+ span.set_status(Status(StatusCode.OK))
673
+ return result
674
+
675
+ except Exception as e:
676
+ span.set_status(Status(StatusCode.ERROR, str(e)))
677
+ span.record_exception(e)
678
+ raise
679
+
680
+ return wrapper
681
+
682
+
683
+ def create_team_internal_async_wrapper(tracer, streaming_context_manager):
684
+ """Create an async wrapper for Team internal methods (_arun) that manages team span lifecycle."""
685
+
686
+ async def wrapper(wrapped, instance, args, kwargs):
687
+ # Get team ID for context storage
688
+ team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
689
+ team_id = str(team_id)
690
+
691
+ # Get team name for span naming
692
+ team_name = getattr(instance, "name", None)
693
+ span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
694
+
695
+ # Check if we already have a team context (from print_response)
696
+ existing_context = streaming_context_manager.get_context(team_id)
697
+
698
+ if existing_context:
699
+ # We're being called from print_response, use existing context
700
+ parent_context, parent_span = existing_context
701
+
702
+ # Execute within the existing team context
703
+ context_token = otel_context.attach(parent_context)
704
+ try:
705
+ with tracer.start_as_current_span(span_name) as span:
706
+ try:
707
+ # Set workflow attributes
708
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
709
+ for key, value in attributes.items():
710
+ span.set_attribute(key, value)
711
+
712
+ # Execute the original function
713
+ result = await wrapped(*args, **kwargs)
714
+
715
+ span.set_status(Status(StatusCode.OK))
716
+ return result
717
+
718
+ except Exception as e:
719
+ span.set_status(Status(StatusCode.ERROR, str(e)))
720
+ span.record_exception(e)
721
+ raise
722
+ finally:
723
+ # Close the parent team span when workflow completes
724
+ if parent_span:
725
+ parent_span.end()
726
+ streaming_context_manager.remove_context(team_id)
727
+ finally:
728
+ otel_context.detach(context_token)
729
+ else:
730
+ # Direct call to _arun, create new team span
731
+ with tracer.start_as_current_span(span_name) as span:
732
+ try:
733
+ # Set workflow attributes
734
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
735
+ for key, value in attributes.items():
736
+ span.set_attribute(key, value)
737
+
738
+ # Execute the original function
739
+ result = await wrapped(*args, **kwargs)
740
+
741
+ span.set_status(Status(StatusCode.OK))
742
+ return result
743
+
744
+ except Exception as e:
745
+ span.set_status(Status(StatusCode.ERROR, str(e)))
746
+ span.record_exception(e)
747
+ raise
748
+
749
+ return wrapper
750
+
751
+
752
+ def create_team_wrapper(tracer, streaming_context_manager):
753
+ """Create a wrapper for Team methods that establishes the team context."""
754
+
755
+ def wrapper(wrapped, instance, args, kwargs):
756
+ # Get team ID for context storage
757
+ team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
758
+ team_id = str(team_id)
759
+
760
+ # Check if streaming is enabled
761
+ is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
762
+
763
+ # Get team name for span naming
764
+ team_name = getattr(instance, "name", None)
765
+ base_span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
766
+
767
+ # For print_response, we need to wrap the internal _run method instead
768
+ # because print_response returns immediately
769
+ if wrapped.__name__ == "print_response":
770
+ # Create team span but don't manage it here
771
+ span = tracer.start_span(base_span_name)
772
+
773
+ try:
774
+ # Set team attributes
775
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
776
+ for key, value in attributes.items():
777
+ span.set_attribute(key, value)
778
+
779
+ # Store context for child spans
780
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
781
+ streaming_context_manager.store_context(team_id, current_context, span)
782
+
783
+ # The span will be closed by the internal _run method
784
+ # Just execute print_response normally
785
+ result = wrapped(*args, **kwargs)
786
+ return result
787
+
788
+ except Exception as e:
789
+ span.set_status(Status(StatusCode.ERROR, str(e)))
790
+ span.record_exception(e)
791
+ span.end()
792
+ streaming_context_manager.remove_context(team_id)
793
+ raise
794
+ else:
795
+ # For run/arun methods, use standard span management
796
+ span = tracer.start_span(base_span_name)
797
+
798
+ try:
799
+ # Set team attributes
800
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
801
+ for key, value in attributes.items():
802
+ span.set_attribute(key, value)
803
+
804
+ # Store context for child spans
805
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
806
+ streaming_context_manager.store_context(team_id, current_context, span)
807
+
808
+ # Execute the original function within team context
809
+ context_token = otel_context.attach(current_context)
810
+ try:
811
+ result = wrapped(*args, **kwargs)
812
+
813
+ # For streaming results, wrap them to keep span alive
814
+ if is_streaming and hasattr(result, "__iter__"):
815
+ return StreamingResultWrapper(result, span, team_id, current_context, streaming_context_manager)
816
+ else:
817
+ # Non-streaming, close span
818
+ span.end()
819
+ streaming_context_manager.remove_context(team_id)
820
+ return result
821
+
822
+ finally:
823
+ otel_context.detach(context_token)
824
+
825
+ except Exception as e:
826
+ span.set_status(Status(StatusCode.ERROR, str(e)))
827
+ span.record_exception(e)
828
+ span.end()
829
+ streaming_context_manager.remove_context(team_id)
830
+ raise
831
+
832
+ return wrapper
833
+
834
+
835
+ def create_team_async_wrapper(tracer, streaming_context_manager):
836
+ """Create an async wrapper for Team methods that establishes the team context."""
837
+
838
+ async def wrapper(wrapped, instance, args, kwargs):
839
+ # Get team ID for context storage
840
+ team_id = getattr(instance, "team_id", None) or getattr(instance, "id", None) or id(instance)
841
+ team_id = str(team_id)
842
+
843
+ # Check if streaming is enabled
844
+ is_streaming = kwargs.get("stream", getattr(instance, "stream", False))
845
+
846
+ # Get team name for span naming
847
+ team_name = getattr(instance, "name", None)
848
+ span_name = f"{team_name}.agno.team.run.workflow" if team_name else "agno.team.run.workflow"
849
+
850
+ # Create team span
851
+ span = tracer.start_span(span_name)
852
+
853
+ try:
854
+ # Set team attributes
855
+ attributes = get_team_run_attributes(args=(instance,) + args, kwargs=kwargs)
856
+ for key, value in attributes.items():
857
+ span.set_attribute(key, value)
858
+
859
+ # Store context for child spans - capture current context with active span
860
+ current_context = trace.set_span_in_context(span, otel_context.get_current())
861
+ streaming_context_manager.store_context(team_id, current_context, span)
862
+
863
+ # Execute the original function within team context
864
+ context_token = otel_context.attach(current_context)
865
+ try:
866
+ result = await wrapped(*args, **kwargs)
867
+
868
+ # For non-streaming, close the span
869
+ if not is_streaming:
870
+ span.end()
871
+ streaming_context_manager.remove_context(team_id)
872
+
873
+ return result
874
+ finally:
875
+ otel_context.detach(context_token)
876
+
877
+ except Exception as e:
878
+ span.set_status(Status(StatusCode.ERROR, str(e)))
879
+ span.record_exception(e)
880
+ span.end()
881
+ streaming_context_manager.remove_context(team_id)
882
+ raise
883
+
884
+ return wrapper
885
+
886
+
887
+ def create_storage_read_wrapper(tracer, streaming_context_manager):
888
+ """Create a wrapper for storage read operations with cache-aware span naming."""
889
+
890
+ def wrapper(wrapped, instance, args, kwargs):
891
+ # Start with a basic span name
892
+ span_name = "agno.workflow.storage.read"
893
+
894
+ with tracer.start_as_current_span(span_name) as span:
895
+ try:
896
+ # Set flag to indicate we're in a storage operation
897
+ SessionStateProxy._set_storage_operation(True)
898
+
899
+ # Set initial attributes
900
+ attributes = get_storage_read_attributes(args=(instance,) + args, kwargs=kwargs)
901
+ for key, value in attributes.items():
902
+ span.set_attribute(key, value)
903
+
904
+ # Execute the original function
905
+ result = wrapped(*args, **kwargs)
906
+
907
+ # Set result attributes including cache hit/miss
908
+ result_attributes = get_storage_read_attributes(
909
+ args=(instance,) + args, kwargs=kwargs, return_value=result
910
+ )
911
+ for key, value in result_attributes.items():
912
+ if key not in attributes: # Avoid duplicates
913
+ span.set_attribute(key, value)
914
+
915
+ # Update span name based on result and cache state
916
+ if hasattr(instance, "session_state") and isinstance(instance.session_state, dict):
917
+ cache_size = len(instance.session_state)
918
+ if result is not None:
919
+ span.update_name(f"Storage.Read.Hit[cache:{cache_size}]")
920
+ else:
921
+ span.update_name(f"Storage.Read.Miss[cache:{cache_size}]")
922
+ else:
923
+ # No cache info available
924
+ if result is not None:
925
+ span.update_name("Storage.Read.Hit")
926
+ else:
927
+ span.update_name("Storage.Read.Miss")
928
+
929
+ span.set_status(Status(StatusCode.OK))
930
+ return result
931
+
932
+ except Exception as e:
933
+ span.set_status(Status(StatusCode.ERROR, str(e)))
934
+ span.record_exception(e)
935
+ raise
936
+ finally:
937
+ # Clear the flag when done
938
+ SessionStateProxy._set_storage_operation(False)
939
+
940
+ return wrapper
941
+
942
+
943
+ def create_storage_write_wrapper(tracer, streaming_context_manager):
944
+ """Create a wrapper for storage write operations with descriptive span naming."""
945
+
946
+ def wrapper(wrapped, instance, args, kwargs):
947
+ # Start with a basic span name
948
+ span_name = "agno.workflow.storage.write"
949
+
950
+ with tracer.start_as_current_span(span_name) as span:
951
+ try:
952
+ # Set flag to indicate we're in a storage operation
953
+ SessionStateProxy._set_storage_operation(True)
954
+
955
+ # Set initial attributes
956
+ attributes = get_storage_write_attributes(args=(instance,) + args, kwargs=kwargs)
957
+ for key, value in attributes.items():
958
+ span.set_attribute(key, value)
959
+
960
+ # Execute the original function
961
+ result = wrapped(*args, **kwargs)
962
+
963
+ # Set result attributes
964
+ result_attributes = get_storage_write_attributes(
965
+ args=(instance,) + args, kwargs=kwargs, return_value=result
966
+ )
967
+ for key, value in result_attributes.items():
968
+ if key not in attributes: # Avoid duplicates
969
+ span.set_attribute(key, value)
970
+
971
+ # Update span name to show cache state after write
972
+ if hasattr(instance, "session_state") and isinstance(instance.session_state, dict):
973
+ cache_size = len(instance.session_state)
974
+ span.update_name(f"Storage.Write[cache:{cache_size}]")
975
+ else:
976
+ span.update_name("Storage.Write")
977
+
978
+ span.set_status(Status(StatusCode.OK))
979
+ return result
980
+
981
+ except Exception as e:
982
+ span.set_status(Status(StatusCode.ERROR, str(e)))
983
+ span.record_exception(e)
984
+ raise
985
+ finally:
986
+ # Clear the flag when done
987
+ SessionStateProxy._set_storage_operation(False)
988
+
989
+ return wrapper
990
+
991
+
992
+ class SessionStateProxy(dict):
993
+ """Proxy class for session_state that instruments cache operations."""
994
+
995
+ # Thread-local storage to track if we're in a storage operation
996
+ _thread_local = threading.local()
997
+
998
+ def __init__(self, original_dict, workflow, tracer):
999
+ super().__init__(original_dict)
1000
+ self._workflow = workflow
1001
+ self._tracer = tracer
1002
+
1003
+ @classmethod
1004
+ def _in_storage_operation(cls):
1005
+ """Check if we're currently in a storage operation."""
1006
+ return getattr(cls._thread_local, "in_storage_operation", False)
1007
+
1008
+ @classmethod
1009
+ def _set_storage_operation(cls, value):
1010
+ """Set whether we're in a storage operation."""
1011
+ cls._thread_local.in_storage_operation = value
1012
+
1013
+ def get(self, key, default=None):
1014
+ """Instrumented get method for cache checking."""
1015
+ # Check if we're already in a storage operation to avoid nested spans
1016
+ if self._in_storage_operation():
1017
+ # We're inside a storage operation, skip instrumentation
1018
+ return super().get(key, default)
1019
+
1020
+ span_name = "Cache.Check"
1021
+
1022
+ with self._tracer.start_as_current_span(span_name) as span:
1023
+ # Set cache attributes
1024
+ span.set_attribute("cache.key", str(key))
1025
+ span.set_attribute("cache.size", len(self))
1026
+ span.set_attribute("cache.keys", json.dumps(list(self.keys())))
1027
+
1028
+ # Get workflow info
1029
+ if hasattr(self._workflow, "workflow_id") and self._workflow.workflow_id:
1030
+ span.set_attribute("cache.workflow_id", str(self._workflow.workflow_id))
1031
+ if hasattr(self._workflow, "session_id") and self._workflow.session_id:
1032
+ span.set_attribute("cache.session_id", str(self._workflow.session_id))
1033
+
1034
+ # Call the original method
1035
+ result = super().get(key, default)
1036
+
1037
+ # Update span based on result
1038
+ if result is not None and result != default:
1039
+ span.set_attribute("cache.hit", True)
1040
+ span.set_attribute("cache.result", "hit")
1041
+ span.update_name(f"Cache.Hit[{len(self)} entries]")
1042
+
1043
+ # Add value info
1044
+ if isinstance(result, str):
1045
+ span.set_attribute("cache.value_size", len(result))
1046
+ if len(result) <= 100:
1047
+ span.set_attribute("cache.value", result)
1048
+ else:
1049
+ span.set_attribute("cache.value_preview", result[:100] + "...")
1050
+ else:
1051
+ span.set_attribute("cache.hit", False)
1052
+ span.set_attribute("cache.result", "miss")
1053
+ span.update_name(f"Cache.Miss[{len(self)} entries]")
1054
+
1055
+ span.set_status(Status(StatusCode.OK))
1056
+ return result
1057
+
1058
+ def __setitem__(self, key, value):
1059
+ """Instrumented setitem method for cache storing."""
1060
+ # Check if we're already in a storage operation to avoid nested spans
1061
+ if self._in_storage_operation():
1062
+ # We're inside a storage operation, skip instrumentation
1063
+ return super().__setitem__(key, value)
1064
+
1065
+ span_name = "Cache.Store"
1066
+
1067
+ with self._tracer.start_as_current_span(span_name) as span:
1068
+ # Set cache attributes
1069
+ span.set_attribute("cache.key", str(key))
1070
+
1071
+ # Get workflow info
1072
+ if hasattr(self._workflow, "workflow_id") and self._workflow.workflow_id:
1073
+ span.set_attribute("cache.workflow_id", str(self._workflow.workflow_id))
1074
+ if hasattr(self._workflow, "session_id") and self._workflow.session_id:
1075
+ span.set_attribute("cache.session_id", str(self._workflow.session_id))
1076
+
1077
+ # Call the original method
1078
+ super().__setitem__(key, value)
1079
+
1080
+ # Set post-store attributes
1081
+ span.set_attribute("cache.size", len(self))
1082
+ span.set_attribute("cache.keys", json.dumps(list(self.keys())))
1083
+
1084
+ # Add value info
1085
+ if isinstance(value, str):
1086
+ span.set_attribute("cache.value_size", len(value))
1087
+ if len(value) <= 100:
1088
+ span.set_attribute("cache.value", value)
1089
+ else:
1090
+ span.set_attribute("cache.value_preview", value[:100] + "...")
1091
+
1092
+ span.update_name(f"Cache.Store[{len(self)} entries]")
1093
+ span.set_status(Status(StatusCode.OK))
1094
+
1095
+
1096
+ def create_workflow_init_wrapper(tracer):
1097
+ """Wrapper to instrument workflow initialization and wrap session_state."""
1098
+
1099
+ def wrapper(wrapped, instance, args, kwargs):
1100
+ # Call the original __init__
1101
+ result = wrapped(*args, **kwargs)
1102
+
1103
+ # Wrap session_state if it exists
1104
+ if hasattr(instance, "session_state") and isinstance(instance.session_state, dict):
1105
+ # Replace session_state with our proxy
1106
+ original_state = instance.session_state
1107
+ instance.session_state = SessionStateProxy(original_state, instance, tracer)
1108
+
1109
+ return result
1110
+
1111
+ return wrapper
1112
+
1113
+
1114
+ def get_agent_context_for_llm():
1115
+ """Helper function for LLM instrumentation to get current agent context."""
1116
+ current_context = otel_context.get_current()
1117
+ current_span = trace.get_current_span(current_context)
1118
+
1119
+ # Check if we're already in an agent span
1120
+ if current_span and hasattr(current_span, "name") and "agent" in current_span.name:
1121
+ return current_context, current_span
1122
+
1123
+ # Try to find stored agent context by checking active contexts
1124
+ # This is a fallback for cases where context isn't properly propagated
1125
+ return None, None
1126
+
1127
+
1128
+ class AgnoInstrumentor(CommonInstrumentor):
1129
+ """Agno instrumentation class."""
1130
+
1131
+ def __init__(self):
1132
+ """Initialize the Agno instrumentor."""
1133
+ self._streaming_context_manager = StreamingContextManager()
1134
+
1135
+ # Create instrumentor config with populated wrapped methods
1136
+ config = InstrumentorConfig(
1137
+ library_name="agentops.instrumentation.agno",
1138
+ library_version="0.1.0",
1139
+ wrapped_methods=self._get_initial_wrapped_methods(),
1140
+ metrics_enabled=True,
1141
+ dependencies=["agno >= 0.1.0"],
1142
+ )
1143
+
1144
+ super().__init__(config)
1145
+
1146
+ def _get_initial_wrapped_methods(self) -> List[WrapConfig]:
1147
+ """Return list of methods to be wrapped during initialization."""
1148
+ # Only return the standard wrapped methods that don't need custom wrappers
1149
+ return WRAPPED_METHODS.copy()
1150
+
1151
+ def _create_metrics(self, meter: Meter) -> Dict[str, Any]:
1152
+ """Create metrics for the instrumentor.
1153
+
1154
+ Returns a dictionary of metric name to metric instance.
1155
+ """
1156
+ # Create standard metrics for LLM operations
1157
+ return StandardMetrics.create_standard_metrics(meter)
1158
+
1159
+ def _initialize(self, **kwargs):
1160
+ """Perform custom initialization."""
1161
+ logger.info("Agno instrumentation: Beginning immediate instrumentation")
1162
+ # Perform wrapping immediately instead of with a delay
1163
+ try:
1164
+ self._perform_wrapping()
1165
+ logger.info("Agno instrumentation: Immediate instrumentation completed successfully")
1166
+ except Exception as e:
1167
+ logger.error(f"Failed to perform immediate wrapping: {e}")
1168
+
1169
+ def _custom_wrap(self, **kwargs):
1170
+ """Skip custom wrapping during initialization - it's done in _initialize."""
1171
+ pass
1172
+
1173
+ def _perform_wrapping(self):
1174
+ """Actually perform the wrapping - called after imports are complete."""
1175
+ if not self._tracer:
1176
+ logger.debug("No tracer available for Agno wrapping")
1177
+ return
1178
+
1179
+ from agentops.instrumentation.common.wrappers import wrap_function_wrapper, WrapConfig, wrap
1180
+
1181
+ # Import Agno modules now that they should be fully loaded
1182
+ try:
1183
+ import agno.agent
1184
+ import agno.workflow.workflow
1185
+ import agno.tools.function
1186
+ import agno.team.team # Noqa: F401
1187
+ except ImportError as e:
1188
+ logger.error(f"Failed to import Agno modules for wrapping: {e}")
1189
+ return
1190
+
1191
+ # First wrap the standard workflow session methods using the standard wrapper
1192
+ session_methods = [
1193
+ WrapConfig(
1194
+ trace_name="agno.workflow.session.load_session",
1195
+ package="agno.workflow.workflow",
1196
+ class_name="Workflow",
1197
+ method_name="load_session",
1198
+ handler=get_workflow_session_attributes,
1199
+ ),
1200
+ WrapConfig(
1201
+ trace_name="agno.workflow.session.new_session",
1202
+ package="agno.workflow.workflow",
1203
+ class_name="Workflow",
1204
+ method_name="new_session",
1205
+ handler=get_workflow_session_attributes,
1206
+ ),
1207
+ # Note: read_from_storage and write_to_storage use custom wrappers below
1208
+ ]
1209
+
1210
+ wrapped_count = 0
1211
+ for wrap_config in session_methods:
1212
+ try:
1213
+ wrap(wrap_config, self._tracer)
1214
+ wrapped_count += 1
1215
+ except Exception as e:
1216
+ logger.debug(f"Failed to wrap {wrap_config}: {e}")
1217
+
1218
+ # Now wrap the streaming methods that need custom wrappers
1219
+ streaming_methods = [
1220
+ # Streaming agent methods
1221
+ ("agno.agent", "Agent.run", self._create_streaming_agent_wrapper()),
1222
+ ("agno.agent", "Agent.arun", self._create_streaming_agent_async_wrapper()),
1223
+ # Streaming workflow methods
1224
+ ("agno.workflow.workflow", "Workflow.run_workflow", self._create_streaming_workflow_wrapper()),
1225
+ ("agno.workflow.workflow", "Workflow.arun_workflow", self._create_streaming_workflow_async_wrapper()),
1226
+ # Streaming tool execution
1227
+ ("agno.tools.function", "FunctionCall.execute", self._create_streaming_tool_wrapper()),
1228
+ # Metrics wrapper
1229
+ ("agno.agent", "Agent._set_session_metrics", self._create_metrics_wrapper()),
1230
+ # Team methods - wrap all public and internal methods
1231
+ ("agno.team.team", "Team.print_response", self._create_team_wrapper()),
1232
+ ("agno.team.team", "Team.run", self._create_team_wrapper()),
1233
+ ("agno.team.team", "Team.arun", self._create_team_async_wrapper()),
1234
+ # Team internal methods with special handling
1235
+ ("agno.team.team", "Team._run", self._create_team_internal_wrapper()),
1236
+ ("agno.team.team", "Team._arun", self._create_team_internal_async_wrapper()),
1237
+ # Storage methods with custom wrappers for cache-aware naming
1238
+ ("agno.workflow.workflow", "Workflow.read_from_storage", self._create_storage_read_wrapper()),
1239
+ ("agno.workflow.workflow", "Workflow.write_to_storage", self._create_storage_write_wrapper()),
1240
+ # Workflow init wrapper to instrument session_state
1241
+ ("agno.workflow.workflow", "Workflow.__init__", self._create_workflow_init_wrapper()),
1242
+ ]
1243
+
1244
+ for package, method, wrapper in streaming_methods:
1245
+ try:
1246
+ wrap_function_wrapper(package, method, wrapper)
1247
+ wrapped_count += 1
1248
+ except Exception as e:
1249
+ logger.debug(f"Failed to wrap {package}.{method}: {e}")
1250
+
1251
+ if wrapped_count > 0:
1252
+ logger.info(f"Agno instrumentation: Successfully wrapped {wrapped_count} methods")
1253
+ else:
1254
+ logger.warning("Agno instrumentation: No methods were successfully wrapped")
1255
+
1256
+ def _custom_unwrap(self, **kwargs):
1257
+ """Perform custom unwrapping."""
1258
+ # Clear streaming contexts
1259
+ self._streaming_context_manager.clear_all()
1260
+ logger.info("Agno instrumentation removed successfully")
1261
+
1262
+ # Method wrappers converted to instance methods
1263
+ def _create_streaming_agent_wrapper(self, args=None, kwargs=None, return_value=None):
1264
+ """Wrapper function for streaming agent methods."""
1265
+ return create_streaming_agent_wrapper(self._tracer, self._streaming_context_manager)
1266
+
1267
+ def _create_streaming_agent_async_wrapper(self, args=None, kwargs=None, return_value=None):
1268
+ """Wrapper function for async streaming agent methods."""
1269
+ return create_streaming_agent_async_wrapper(self._tracer, self._streaming_context_manager)
1270
+
1271
+ def _create_streaming_workflow_wrapper(self, args=None, kwargs=None, return_value=None):
1272
+ """Wrapper function for streaming workflow methods."""
1273
+ return create_streaming_workflow_wrapper(self._tracer, self._streaming_context_manager)
1274
+
1275
+ def _create_streaming_workflow_async_wrapper(self, args=None, kwargs=None, return_value=None):
1276
+ """Wrapper function for async streaming workflow methods."""
1277
+ return create_streaming_workflow_async_wrapper(self._tracer, self._streaming_context_manager)
1278
+
1279
+ def _create_streaming_tool_wrapper(self, args=None, kwargs=None, return_value=None):
1280
+ """Wrapper function for streaming tool methods."""
1281
+ return create_streaming_tool_wrapper(self._tracer, self._streaming_context_manager)
1282
+
1283
+ def _create_metrics_wrapper(self, args=None, kwargs=None, return_value=None):
1284
+ """Wrapper function for metrics methods."""
1285
+ return create_metrics_wrapper(self._tracer, self._streaming_context_manager)
1286
+
1287
+ def _create_team_wrapper(self, args=None, kwargs=None, return_value=None):
1288
+ """Wrapper function for team methods."""
1289
+ return create_team_wrapper(self._tracer, self._streaming_context_manager)
1290
+
1291
+ def _create_team_async_wrapper(self, args=None, kwargs=None, return_value=None):
1292
+ """Wrapper function for async team methods."""
1293
+ return create_team_async_wrapper(self._tracer, self._streaming_context_manager)
1294
+
1295
+ def _create_team_internal_wrapper(self, args=None, kwargs=None, return_value=None):
1296
+ """Wrapper function for team internal methods."""
1297
+ return create_team_internal_wrapper(self._tracer, self._streaming_context_manager)
1298
+
1299
+ def _create_team_internal_async_wrapper(self, args=None, kwargs=None, return_value=None):
1300
+ """Wrapper function for async team internal methods."""
1301
+ return create_team_internal_async_wrapper(self._tracer, self._streaming_context_manager)
1302
+
1303
+ def _create_storage_read_wrapper(self, args=None, kwargs=None, return_value=None):
1304
+ """Wrapper function for storage read operations."""
1305
+ return create_storage_read_wrapper(self._tracer, self._streaming_context_manager)
1306
+
1307
+ def _create_storage_write_wrapper(self, args=None, kwargs=None, return_value=None):
1308
+ """Wrapper function for storage write operations."""
1309
+ return create_storage_write_wrapper(self._tracer, self._streaming_context_manager)
1310
+
1311
+ def _create_workflow_init_wrapper(self, args=None, kwargs=None, return_value=None):
1312
+ """Wrapper function for workflow initialization to instrument session_state."""
1313
+ return create_workflow_init_wrapper(self._tracer)