Skip to content

Commit fd449ad

Browse files
IlumCIIlumCI
authored
Implement telemetry features in agent.py
Add telemetry support for agent execution and LLM calls.
1 parent 44cf4d3 commit fd449ad

1 file changed

Lines changed: 234 additions & 3 deletions

File tree

‎swarms/structs/agent.py‎

Lines changed: 234 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,19 @@
7575
handle_transforms,
7676
)
7777
from swarms.telemetry.main import log_agent_data
78+
79+
try:
80+
from swarms.telemetry.opentelemetry_integration import (
81+
trace_span,
82+
record_metric,
83+
log_event,
84+
)
85+
_OTEL_AVAILABLE = True
86+
except ImportError:
87+
_OTEL_AVAILABLE = False
88+
trace_span = None
89+
record_metric = lambda *args, **kwargs: None
90+
log_event = lambda *args, **kwargs: None
7891
from swarms.tools.base_tool import BaseTool
7992
from swarms.tools.mcp_client_tools import (
8093
execute_multiple_tools_on_multiple_mcp_servers_sync,
@@ -470,6 +483,7 @@ def __init__(
470483
mode: Literal["interactive", "fast", "standard"] = "standard",
471484
publish_to_marketplace: bool = False,
472485
use_cases: Optional[List[Dict[str, Any]]] = None,
486+
enable_telemetry: Optional[bool] = None,
473487
*args,
474488
**kwargs,
475489
):
@@ -623,6 +637,15 @@ def __init__(
623637
self.mode = mode
624638
self.publish_to_marketplace = publish_to_marketplace
625639

640+
if enable_telemetry is None:
641+
import os
642+
self.enable_telemetry = (
643+
_OTEL_AVAILABLE
644+
and os.getenv("OTEL_ENABLED", "true").lower() == "true"
645+
)
646+
else:
647+
self.enable_telemetry = enable_telemetry and _OTEL_AVAILABLE
648+
626649
# Initialize transforms
627650
if transforms is None:
628651
self.transforms = None
@@ -1215,7 +1238,42 @@ def _run(
12151238
agent(task="Summarize this document.", img="path/to/image.jpg")
12161239
agent(task="Analyze this image.", img="path/to/image.jpg", is_last=True)
12171240
"""
1241+
span_attributes = {}
1242+
if self.enable_telemetry and trace_span:
1243+
span_attributes = {
1244+
"agent.id": self.id,
1245+
"agent.name": self.agent_name,
1246+
"agent.model": self.get_current_model(),
1247+
"agent.max_loops": str(self.max_loops),
1248+
"task.length": len(str(task)) if task else 0,
1249+
}
1250+
1251+
span_context = (
1252+
trace_span(
1253+
f"agent.run.{self.agent_name}",
1254+
attributes=span_attributes,
1255+
)
1256+
if (self.enable_telemetry and trace_span)
1257+
else None
1258+
)
1259+
1260+
if span_context:
1261+
span_manager = span_context.__enter__()
1262+
else:
1263+
from contextlib import nullcontext
1264+
span_manager = nullcontext()
1265+
12181266
try:
1267+
if self.enable_telemetry and record_metric:
1268+
record_metric(
1269+
"agent.executions.total",
1270+
1,
1271+
{
1272+
"agent_name": self.agent_name,
1273+
"model": self.get_current_model(),
1274+
},
1275+
metric_type="counter",
1276+
)
12191277

12201278
self.check_if_no_prompt_then_autogenerate(task)
12211279

@@ -1250,6 +1308,17 @@ def _run(
12501308
):
12511309
loop_count += 1
12521310

1311+
if self.enable_telemetry and record_metric:
1312+
record_metric(
1313+
"agent.loops",
1314+
1,
1315+
{
1316+
"agent_name": self.agent_name,
1317+
"loop_number": str(loop_count),
1318+
},
1319+
metric_type="counter",
1320+
)
1321+
12531322
# Handle RAG query every loop
12541323
if (
12551324
self.long_term_memory is not None
@@ -1457,15 +1526,63 @@ def _run(
14571526

14581527
self.save()
14591528

1460-
# Output formatting based on output_type
1461-
return history_output_formatter(
1529+
if self.enable_telemetry and record_metric:
1530+
record_metric(
1531+
"agent.executions.success",
1532+
1,
1533+
{"agent_name": self.agent_name},
1534+
metric_type="counter",
1535+
)
1536+
1537+
result = history_output_formatter(
14621538
self.short_memory, type=self.output_type
14631539
)
14641540

1541+
if span_context:
1542+
try:
1543+
span_context.__exit__(None, None, None)
1544+
except Exception:
1545+
pass
1546+
1547+
return result
1548+
14651549
except Exception as error:
1550+
if self.enable_telemetry:
1551+
if record_metric:
1552+
record_metric(
1553+
"agent.executions.errors",
1554+
1,
1555+
{
1556+
"agent_name": self.agent_name,
1557+
"error_type": type(error).__name__,
1558+
},
1559+
metric_type="counter",
1560+
)
1561+
if log_event:
1562+
log_event(
1563+
f"Agent {self.agent_name} execution failed",
1564+
level="ERROR",
1565+
attributes={
1566+
"agent_id": self.id,
1567+
"error_type": type(error).__name__,
1568+
"error_message": str(error)[:200],
1569+
},
1570+
)
1571+
1572+
if span_context:
1573+
try:
1574+
span_context.__exit__(type(error), error, None)
1575+
except Exception:
1576+
pass
1577+
14661578
self._handle_run_error(error)
14671579

14681580
except KeyboardInterrupt as error:
1581+
if span_context:
1582+
try:
1583+
span_context.__exit__(type(error), error, None)
1584+
except Exception:
1585+
pass
14691586
self._handle_run_error(error)
14701587

14711588
def _handle_run_error(self, error: any):
@@ -2565,6 +2682,33 @@ def call_llm(
25652682
if "is_last" in kwargs:
25662683
del kwargs["is_last"]
25672684

2685+
span_attrs = {}
2686+
if self.enable_telemetry and trace_span:
2687+
span_attrs = {
2688+
"agent.name": self.agent_name,
2689+
"agent.model": self.get_current_model(),
2690+
"agent.loop": str(current_loop),
2691+
"task.length": len(task),
2692+
"has_image": str(img is not None),
2693+
}
2694+
2695+
llm_span_context = (
2696+
trace_span(
2697+
"agent.llm.call",
2698+
attributes=span_attrs,
2699+
)
2700+
if (self.enable_telemetry and trace_span)
2701+
else None
2702+
)
2703+
2704+
if llm_span_context:
2705+
llm_span_manager = llm_span_context.__enter__()
2706+
else:
2707+
from contextlib import nullcontext
2708+
llm_span_manager = nullcontext()
2709+
2710+
start_time = time.time()
2711+
25682712
try:
25692713
# Set streaming parameter in LLM if streaming is enabled
25702714
if self.streaming_on and hasattr(self.llm, "stream"):
@@ -2640,11 +2784,57 @@ def on_chunk_received(chunk: str):
26402784
# Restore original stream setting
26412785
self.llm.stream = original_stream
26422786

2643-
# Return the complete response for further processing
2787+
if self.enable_telemetry and record_metric:
2788+
execution_time = time.time() - start_time
2789+
record_metric(
2790+
"agent.llm.call.duration",
2791+
execution_time,
2792+
{
2793+
"agent_name": self.agent_name,
2794+
"model": self.get_current_model(),
2795+
},
2796+
)
2797+
record_metric(
2798+
"agent.llm.calls.total",
2799+
1,
2800+
{"model": self.get_current_model()},
2801+
metric_type="counter",
2802+
)
2803+
2804+
if llm_span_context:
2805+
try:
2806+
llm_span_context.__exit__(None, None, None)
2807+
except Exception:
2808+
pass
2809+
26442810
return complete_response
26452811
else:
26462812
# Restore original stream setting
26472813
self.llm.stream = original_stream
2814+
2815+
if self.enable_telemetry and record_metric:
2816+
execution_time = time.time() - start_time
2817+
record_metric(
2818+
"agent.llm.call.duration",
2819+
execution_time,
2820+
{
2821+
"agent_name": self.agent_name,
2822+
"model": self.get_current_model(),
2823+
},
2824+
)
2825+
record_metric(
2826+
"agent.llm.calls.total",
2827+
1,
2828+
{"model": self.get_current_model()},
2829+
metric_type="counter",
2830+
)
2831+
2832+
if llm_span_context:
2833+
try:
2834+
llm_span_context.__exit__(None, None, None)
2835+
except Exception:
2836+
pass
2837+
26482838
return streaming_response
26492839
else:
26502840
args = {
@@ -2656,6 +2846,29 @@ def on_chunk_received(chunk: str):
26562846

26572847
out = self.llm.run(**args, **kwargs)
26582848

2849+
if self.enable_telemetry and record_metric:
2850+
execution_time = time.time() - start_time
2851+
record_metric(
2852+
"agent.llm.call.duration",
2853+
execution_time,
2854+
{
2855+
"agent_name": self.agent_name,
2856+
"model": self.get_current_model(),
2857+
},
2858+
)
2859+
record_metric(
2860+
"agent.llm.calls.total",
2861+
1,
2862+
{"model": self.get_current_model()},
2863+
metric_type="counter",
2864+
)
2865+
2866+
if llm_span_context:
2867+
try:
2868+
llm_span_context.__exit__(None, None, None)
2869+
except Exception:
2870+
pass
2871+
26592872
return out
26602873

26612874
except (
@@ -2665,6 +2878,24 @@ def on_chunk_received(chunk: str):
26652878
AuthenticationError,
26662879
Exception,
26672880
) as e:
2881+
if self.enable_telemetry:
2882+
if record_metric:
2883+
record_metric(
2884+
"agent.llm.call.errors",
2885+
1,
2886+
{
2887+
"agent_name": self.agent_name,
2888+
"model": self.get_current_model(),
2889+
"error_type": type(e).__name__,
2890+
},
2891+
metric_type="counter",
2892+
)
2893+
if llm_span_context:
2894+
try:
2895+
llm_span_context.__exit__(type(e), e, None)
2896+
except Exception:
2897+
pass
2898+
26682899
logger.error(
26692900
f"Error calling LLM with model '{self.get_current_model()}': {e}. "
26702901
f"Task: {task}, Args: {args}, Kwargs: {kwargs} Traceback: {traceback.format_exc()}"

0 commit comments

Comments
 (0)