phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-23 08:42:28 +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

91
llama_stack/distribution/routing_tables/toolgroups.py
View file

@ -43,6 +43,8 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
return super().get_provider_impl(routing_key, provider_id)
async def list_tools(self, toolgroup_id: str | None = None) -> ListToolsResponse:
logger.debug(f"Listing tools for toolgroup_id: {toolgroup_id}")
if toolgroup_id:
if group_id := parse_toolgroup_from_toolgroup_name_pair(toolgroup_id):
toolgroup_id = group_id
@ -53,33 +55,59 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
all_tools = []
for toolgroup in toolgroups:
if toolgroup.identifier not in self.toolgroups_to_tools:
logger.debug(f"Toolgroup {toolgroup.identifier} not in cache, indexing...")
await self._index_tools(toolgroup)
all_tools.extend(self.toolgroups_to_tools[toolgroup.identifier])
cached_tools = self.toolgroups_to_tools.get(toolgroup.identifier, [])
logger.debug(f"Found {len(cached_tools)} cached tools for toolgroup {toolgroup.identifier}")
all_tools.extend(cached_tools)
logger.debug(f"Returning {len(all_tools)} total tools")
return ListToolsResponse(data=all_tools)
async def _index_tools(self, toolgroup: ToolGroup):
provider_impl = super().get_provider_impl(toolgroup.identifier, toolgroup.provider_id)
tooldefs_response = await provider_impl.list_runtime_tools(toolgroup.identifier, toolgroup.mcp_endpoint)
try:
provider_impl = super().get_provider_impl(toolgroup.identifier, toolgroup.provider_id)
logger.debug(f"Indexing tools for toolgroup {toolgroup.identifier} with provider {toolgroup.provider_id}")
# TODO: kill this Tool vs ToolDef distinction
tooldefs = tooldefs_response.data
tools = []
for t in tooldefs:
tools.append(
Tool(
identifier=t.name,
toolgroup_id=toolgroup.identifier,
description=t.description or "",
parameters=t.parameters or [],
metadata=t.metadata,
provider_id=toolgroup.provider_id,
if toolgroup.mcp_endpoint:
logger.debug(f"Toolgroup {toolgroup.identifier} has MCP endpoint: {toolgroup.mcp_endpoint.uri}")
tooldefs_response = await provider_impl.list_runtime_tools(toolgroup.identifier, toolgroup.mcp_endpoint)
# TODO: kill this Tool vs ToolDef distinction
tooldefs = tooldefs_response.data
tools = []
for t in tooldefs:
tools.append(
Tool(
identifier=t.name,
toolgroup_id=toolgroup.identifier,
description=t.description or "",
parameters=t.parameters or [],
metadata=t.metadata,
provider_id=toolgroup.provider_id,
)
)
)
self.toolgroups_to_tools[toolgroup.identifier] = tools
for tool in tools:
self.tool_to_toolgroup[tool.identifier] = toolgroup.identifier
self.toolgroups_to_tools[toolgroup.identifier] = tools
for tool in tools:
self.tool_to_toolgroup[tool.identifier] = toolgroup.identifier
logger.info(f"Successfully indexed {len(tools)} tools for toolgroup {toolgroup.identifier}")
except Exception as e:
logger.warning(f"Failed to index tools for toolgroup {toolgroup.identifier}: {e}")
# Don't let tool indexing failures crash the system
# Initialize empty tools list so the toolgroup still exists
self.toolgroups_to_tools[toolgroup.identifier] = []
if toolgroup.mcp_endpoint:
logger.info(
f"Toolgroup {toolgroup.identifier} has MCP endpoint - tools may be available after authentication"
)
else:
logger.error(f"Non-MCP toolgroup {toolgroup.identifier} failed to index tools: {e}")
# Don't raise - we want the system to continue running even if tool indexing fails
async def list_tool_groups(self) -> ListToolGroupsResponse:
return ListToolGroupsResponse(data=await self.get_all_with_type("tool_group"))
@ -119,7 +147,12 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
# the tools should first list the tools and then use them. but there are assumptions
# baked in some of the code and tests right now.
if not toolgroup.mcp_endpoint:
await self._index_tools(toolgroup)
try:
await self._index_tools(toolgroup)
except Exception as e:
logger.error(f"Failed to index tools during toolgroup registration for {toolgroup_id}: {e}")
# Don't fail the registration - the toolgroup can still be used
# Tools may become available later
return toolgroup
async def unregister_toolgroup(self, toolgroup_id: str) -> None:
@ -128,5 +161,23 @@ class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
raise ValueError(f"Tool group {toolgroup_id} not found")
await self.unregister_object(tool_group)
async def refresh_tools(self, toolgroup_id: str) -> None:
"""Refresh tools for a specific toolgroup, useful for re-indexing after auth becomes available."""
try:
toolgroup = await self.get_tool_group(toolgroup_id)
# Clear existing tools for this toolgroup
if toolgroup_id in self.toolgroups_to_tools:
old_tools = self.toolgroups_to_tools[toolgroup_id]
for tool in old_tools:
if tool.identifier in self.tool_to_toolgroup:
del self.tool_to_toolgroup[tool.identifier]
del self.toolgroups_to_tools[toolgroup_id]
# Re-index tools for this toolgroup
await self._index_tools(toolgroup)
except Exception as e:
# Log error but don't fail - tools may become available later
logger.warning(f"Failed to refresh tools for toolgroup {toolgroup_id}: {e}")
async def shutdown(self) -> None:
pass