Skip to content

Commit 26fd5db

Browse files
authored
fix: add traces for tool calls and mcp tool listing (#3722)
# What does this PR do? Adds traces around tool execution and mcp tool listing for better observability. Closes #3108 ## Test Plan Manually examined traces in jaeger to verify the added information was available. Signed-off-by: Gordon Sim <gsim@redhat.com>
1 parent 4b9ebbf commit 26fd5db

File tree

2 files changed

+40
-19
lines changed

2 files changed

+40
-19
lines changed

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

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
OpenAIMessageParam,
4747
)
4848
from llama_stack.log import get_logger
49+
from llama_stack.providers.utils.telemetry import tracing
4950

5051
from .types import ChatCompletionContext, ChatCompletionResult
5152
from .utils import convert_chat_choice_to_response_message, is_function_tool_call
@@ -595,14 +596,22 @@ async def _process_mcp_tool(
595596
never_allowed = mcp_tool.allowed_tools.never
596597

597598
# Call list_mcp_tools
598-
tool_defs = await list_mcp_tools(
599-
endpoint=mcp_tool.server_url,
600-
headers=mcp_tool.headers or {},
601-
)
599+
tool_defs = None
600+
list_id = f"mcp_list_{uuid.uuid4()}"
601+
attributes = {
602+
"server_label": mcp_tool.server_label,
603+
"server_url": mcp_tool.server_url,
604+
"mcp_list_tools_id": list_id,
605+
}
606+
async with tracing.span("list_mcp_tools", attributes):
607+
tool_defs = await list_mcp_tools(
608+
endpoint=mcp_tool.server_url,
609+
headers=mcp_tool.headers or {},
610+
)
602611

603612
# Create the MCP list tools message
604613
mcp_list_message = OpenAIResponseOutputMessageMCPListTools(
605-
id=f"mcp_list_{uuid.uuid4()}",
614+
id=list_id,
606615
server_label=mcp_tool.server_label,
607616
tools=[],
608617
)

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

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from llama_stack.apis.tools import ToolGroups, ToolInvocationResult, ToolRuntime
3636
from llama_stack.apis.vector_io import VectorIO
3737
from llama_stack.log import get_logger
38+
from llama_stack.providers.utils.telemetry import tracing
3839

3940
from .types import ChatCompletionContext, ToolExecutionResult
4041

@@ -251,12 +252,18 @@ async def _execute_tool(
251252
from llama_stack.providers.utils.tools.mcp import invoke_mcp_tool
252253

253254
mcp_tool = mcp_tool_to_server[function_name]
254-
result = await invoke_mcp_tool(
255-
endpoint=mcp_tool.server_url,
256-
headers=mcp_tool.headers or {},
257-
tool_name=function_name,
258-
kwargs=tool_kwargs,
259-
)
255+
attributes = {
256+
"server_label": mcp_tool.server_label,
257+
"server_url": mcp_tool.server_url,
258+
"tool_name": function_name,
259+
}
260+
async with tracing.span("invoke_mcp_tool", attributes):
261+
result = await invoke_mcp_tool(
262+
endpoint=mcp_tool.server_url,
263+
headers=mcp_tool.headers or {},
264+
tool_name=function_name,
265+
kwargs=tool_kwargs,
266+
)
260267
elif function_name == "knowledge_search":
261268
response_file_search_tool = next(
262269
(t for t in ctx.response_tools if isinstance(t, OpenAIResponseInputToolFileSearch)),
@@ -266,15 +273,20 @@ async def _execute_tool(
266273
# Use vector_stores.search API instead of knowledge_search tool
267274
# to support filters and ranking_options
268275
query = tool_kwargs.get("query", "")
269-
result = await self._execute_knowledge_search_via_vector_store(
270-
query=query,
271-
response_file_search_tool=response_file_search_tool,
272-
)
276+
async with tracing.span("knowledge_search", {}):
277+
result = await self._execute_knowledge_search_via_vector_store(
278+
query=query,
279+
response_file_search_tool=response_file_search_tool,
280+
)
273281
else:
274-
result = await self.tool_runtime_api.invoke_tool(
275-
tool_name=function_name,
276-
kwargs=tool_kwargs,
277-
)
282+
attributes = {
283+
"tool_name": function_name,
284+
}
285+
async with tracing.span("invoke_tool", attributes):
286+
result = await self.tool_runtime_api.invoke_tool(
287+
tool_name=function_name,
288+
kwargs=tool_kwargs,
289+
)
278290
except Exception as e:
279291
error_exc = e
280292

0 commit comments

Comments
 (0)