phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-25 19:38:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

tool index refreshing when token is provided to create_turn

Browse source
This commit is contained in:
Lance Galletti 2025-07-02 15:08:10 -04:00
parent 57745101be
commit 76f593143b
7 changed files with 217 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

34
llama_stack/providers/inline/agents/meta_reference/agent_instance.py
View file

@ -185,6 +185,9 @@ class ChatAgent(ShieldRunnerMixin):
if self.agent_config.name:
span.set_attribute("agent_name", self.agent_config.name)
# Refresh tools for MCP toolgroups in case auth is now available
await self._refresh_mcp_tools(request.toolgroups)
await self._initialize_tools(request.toolgroups)
async for chunk in self._run_turn(request, turn_id):
yield chunk
@ -762,6 +765,35 @@ class ChatAgent(ShieldRunnerMixin):
yield client_message
return
async def _refresh_mcp_tools(
self,
toolgroups_for_turn: list[AgentToolGroup] | None = None,
) -> None:
"""Refresh MCP tools in case authentication is now available."""
# Determine which tools to refresh
tool_groups_to_refresh = toolgroups_for_turn or self.agent_config.toolgroups or []
logger.debug(f"Refreshing MCP tools for {len(tool_groups_to_refresh)} toolgroups")
for toolgroup in tool_groups_to_refresh:
name = toolgroup.name if isinstance(toolgroup, AgentToolGroupWithArgs) else toolgroup
toolgroup_name, _ = self._parse_toolgroup_name(name)
# Only refresh MCP toolgroups (those with endpoints)
try:
tool_group = await self.tool_groups_api.get_tool_group(toolgroup_name)
if tool_group.mcp_endpoint:
logger.debug(
f"Refreshing tools for MCP toolgroup {toolgroup_name} with endpoint {tool_group.mcp_endpoint.uri}"
)
# Refresh tools for this MCP toolgroup
await self.tool_groups_api.refresh_tools(toolgroup_name)
else:
logger.debug(f"Toolgroup {toolgroup_name} is not an MCP toolgroup, skipping refresh")
except Exception as e:
# Log but don't fail - tools may become available later
logger.warning(f"Failed to refresh MCP tools for toolgroup {toolgroup_name}: {e}")
async def _initialize_tools(
self,
toolgroups_for_turn: list[AgentToolGroup] | None = None,
@ -800,6 +832,7 @@ class ChatAgent(ShieldRunnerMixin):
)
for param in tool_def.parameters
},
toolgroup_name="client_tools",
)
for toolgroup_name_with_maybe_tool_name in agent_config_toolgroups:
toolgroup_name, input_tool_name = self._parse_toolgroup_name(toolgroup_name_with_maybe_tool_name)
@ -843,6 +876,7 @@ class ChatAgent(ShieldRunnerMixin):
)
for param in tool_def.parameters
},
toolgroup_name=toolgroup_name,
)
tool_name_to_args[tool_def.identifier] = toolgroup_to_args.get(toolgroup_name, {})