Skip to content

Commit 6389bf5

Browse files
authored
fix: make telemetry optional for agents (#3705)
# What does this PR do? there is a lot of code in the agents API using the telemetry API and its helpers without checking if that API is even enabled. This is the only API besides inference actively using telemetry code, so after this telemetry can be optional for the entire stack resolves #3665 ## Test Plan existing agent tests. Signed-off-by: Charlie Doern <[email protected]>
1 parent e892a3f commit 6389bf5

File tree

4 files changed

+57
-36
lines changed

4 files changed

+57
-36
lines changed

llama_stack/providers/inline/agents/meta_reference/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ async def get_provider_impl(config: MetaReferenceAgentsImplConfig, deps: dict[Ap
2222
deps[Api.tool_runtime],
2323
deps[Api.tool_groups],
2424
policy,
25+
Api.telemetry in deps,
2526
)
2627
await impl.initialize()
2728
return impl

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

Lines changed: 50 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ def __init__(
110110
persistence_store: KVStore,
111111
created_at: str,
112112
policy: list[AccessRule],
113+
telemetry_enabled: bool = False,
113114
):
114115
self.agent_id = agent_id
115116
self.agent_config = agent_config
@@ -120,6 +121,7 @@ def __init__(
120121
self.tool_runtime_api = tool_runtime_api
121122
self.tool_groups_api = tool_groups_api
122123
self.created_at = created_at
124+
self.telemetry_enabled = telemetry_enabled
123125

124126
ShieldRunnerMixin.__init__(
125127
self,
@@ -188,28 +190,30 @@ async def get_messages_from_turns(self, turns: list[Turn]) -> list[Message]:
188190

189191
async def create_and_execute_turn(self, request: AgentTurnCreateRequest) -> AsyncGenerator:
190192
turn_id = str(uuid.uuid4())
191-
span = tracing.get_current_span()
192-
if span:
193-
span.set_attribute("session_id", request.session_id)
194-
span.set_attribute("agent_id", self.agent_id)
195-
span.set_attribute("request", request.model_dump_json())
196-
span.set_attribute("turn_id", turn_id)
197-
if self.agent_config.name:
198-
span.set_attribute("agent_name", self.agent_config.name)
193+
if self.telemetry_enabled:
194+
span = tracing.get_current_span()
195+
if span is not None:
196+
span.set_attribute("session_id", request.session_id)
197+
span.set_attribute("agent_id", self.agent_id)
198+
span.set_attribute("request", request.model_dump_json())
199+
span.set_attribute("turn_id", turn_id)
200+
if self.agent_config.name:
201+
span.set_attribute("agent_name", self.agent_config.name)
199202

200203
await self._initialize_tools(request.toolgroups)
201204
async for chunk in self._run_turn(request, turn_id):
202205
yield chunk
203206

204207
async def resume_turn(self, request: AgentTurnResumeRequest) -> AsyncGenerator:
205-
span = tracing.get_current_span()
206-
if span:
207-
span.set_attribute("agent_id", self.agent_id)
208-
span.set_attribute("session_id", request.session_id)
209-
span.set_attribute("request", request.model_dump_json())
210-
span.set_attribute("turn_id", request.turn_id)
211-
if self.agent_config.name:
212-
span.set_attribute("agent_name", self.agent_config.name)
208+
if self.telemetry_enabled:
209+
span = tracing.get_current_span()
210+
if span is not None:
211+
span.set_attribute("agent_id", self.agent_id)
212+
span.set_attribute("session_id", request.session_id)
213+
span.set_attribute("request", request.model_dump_json())
214+
span.set_attribute("turn_id", request.turn_id)
215+
if self.agent_config.name:
216+
span.set_attribute("agent_name", self.agent_config.name)
213217

214218
await self._initialize_tools()
215219
async for chunk in self._run_turn(request):
@@ -395,9 +399,12 @@ async def run_multiple_shields_wrapper(
395399
touchpoint: str,
396400
) -> AsyncGenerator:
397401
async with tracing.span("run_shields") as span:
398-
span.set_attribute("input", [m.model_dump_json() for m in messages])
402+
if self.telemetry_enabled and span is not None:
403+
span.set_attribute("input", [m.model_dump_json() for m in messages])
404+
if len(shields) == 0:
405+
span.set_attribute("output", "no shields")
406+
399407
if len(shields) == 0:
400-
span.set_attribute("output", "no shields")
401408
return
402409

403410
step_id = str(uuid.uuid4())
@@ -430,7 +437,8 @@ async def run_multiple_shields_wrapper(
430437
)
431438
)
432439
)
433-
span.set_attribute("output", e.violation.model_dump_json())
440+
if self.telemetry_enabled and span is not None:
441+
span.set_attribute("output", e.violation.model_dump_json())
434442

435443
yield CompletionMessage(
436444
content=str(e),
@@ -453,7 +461,8 @@ async def run_multiple_shields_wrapper(
453461
)
454462
)
455463
)
456-
span.set_attribute("output", "no violations")
464+
if self.telemetry_enabled and span is not None:
465+
span.set_attribute("output", "no violations")
457466

458467
async def _run(
459468
self,
@@ -518,8 +527,9 @@ async def _run(
518527
stop_reason: StopReason | None = None
519528

520529
async with tracing.span("inference") as span:
521-
if self.agent_config.name:
522-
span.set_attribute("agent_name", self.agent_config.name)
530+
if self.telemetry_enabled and span is not None:
531+
if self.agent_config.name:
532+
span.set_attribute("agent_name", self.agent_config.name)
523533

524534
def _serialize_nested(value):
525535
"""Recursively serialize nested Pydantic models to dicts."""
@@ -637,18 +647,19 @@ def _add_type(openai_msg: dict) -> OpenAIMessageParam:
637647
else:
638648
raise ValueError(f"Unexpected delta type {type(delta)}")
639649

640-
span.set_attribute("stop_reason", stop_reason or StopReason.end_of_turn)
641-
span.set_attribute(
642-
"input",
643-
json.dumps([json.loads(m.model_dump_json()) for m in input_messages]),
644-
)
645-
output_attr = json.dumps(
646-
{
647-
"content": content,
648-
"tool_calls": [json.loads(t.model_dump_json()) for t in tool_calls],
649-
}
650-
)
651-
span.set_attribute("output", output_attr)
650+
if self.telemetry_enabled and span is not None:
651+
span.set_attribute("stop_reason", stop_reason or StopReason.end_of_turn)
652+
span.set_attribute(
653+
"input",
654+
json.dumps([json.loads(m.model_dump_json()) for m in input_messages]),
655+
)
656+
output_attr = json.dumps(
657+
{
658+
"content": content,
659+
"tool_calls": [json.loads(t.model_dump_json()) for t in tool_calls],
660+
}
661+
)
662+
span.set_attribute("output", output_attr)
652663

653664
n_iter += 1
654665
await self.storage.set_num_infer_iters_in_turn(session_id, turn_id, n_iter)
@@ -756,7 +767,9 @@ def _add_type(openai_msg: dict) -> OpenAIMessageParam:
756767
{
757768
"tool_name": tool_call.tool_name,
758769
"input": message.model_dump_json(),
759-
},
770+
}
771+
if self.telemetry_enabled
772+
else {},
760773
) as span:
761774
tool_execution_start_time = datetime.now(UTC).isoformat()
762775
tool_result = await self.execute_tool_call_maybe(
@@ -771,7 +784,8 @@ def _add_type(openai_msg: dict) -> OpenAIMessageParam:
771784
call_id=tool_call.call_id,
772785
content=tool_result.content,
773786
)
774-
span.set_attribute("output", result_message.model_dump_json())
787+
if self.telemetry_enabled and span is not None:
788+
span.set_attribute("output", result_message.model_dump_json())
775789

776790
# Store tool execution step
777791
tool_execution_step = ToolExecutionStep(

llama_stack/providers/inline/agents/meta_reference/agents.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,13 +64,15 @@ def __init__(
6464
tool_runtime_api: ToolRuntime,
6565
tool_groups_api: ToolGroups,
6666
policy: list[AccessRule],
67+
telemetry_enabled: bool = False,
6768
):
6869
self.config = config
6970
self.inference_api = inference_api
7071
self.vector_io_api = vector_io_api
7172
self.safety_api = safety_api
7273
self.tool_runtime_api = tool_runtime_api
7374
self.tool_groups_api = tool_groups_api
75+
self.telemetry_enabled = telemetry_enabled
7476

7577
self.in_memory_store = InmemoryKVStoreImpl()
7678
self.openai_responses_impl: OpenAIResponsesImpl | None = None
@@ -135,6 +137,7 @@ async def _get_agent_impl(self, agent_id: str) -> ChatAgent:
135137
),
136138
created_at=agent_info.created_at,
137139
policy=self.policy,
140+
telemetry_enabled=self.telemetry_enabled,
138141
)
139142

140143
async def create_agent_session(

llama_stack/providers/registry/agents.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@ def available_providers() -> list[ProviderSpec]:
3636
Api.tool_runtime,
3737
Api.tool_groups,
3838
],
39+
optional_api_dependencies=[
40+
Api.telemetry,
41+
],
3942
description="Meta's reference implementation of an agent system that can use tools, access vector databases, and perform complex reasoning tasks.",
4043
),
4144
]

0 commit comments

Comments
 (0)