chore: split routing_tables into individual files (#2259)

llama_stack/distribution/routing_tables/toolgroups.py

@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.tools import ListToolGroupsResponse, ListToolsResponse, Tool, ToolGroup, ToolGroups, ToolHost
+from llama_stack.distribution.datatypes import (
+    ToolGroupWithACL,
+    ToolWithACL,
+)
+from llama_stack.log import get_logger
+
+from .common import CommonRoutingTableImpl
+
+logger = get_logger(name=__name__, category="core")
+
+
+class ToolGroupsRoutingTable(CommonRoutingTableImpl, ToolGroups):
+    async def list_tools(self, toolgroup_id: str | None = None) -> ListToolsResponse:
+        tools = await self.get_all_with_type("tool")
+        if toolgroup_id:
+            tools = [tool for tool in tools if tool.toolgroup_id == toolgroup_id]
+        return ListToolsResponse(data=tools)
+
+    async def list_tool_groups(self) -> ListToolGroupsResponse:
+        return ListToolGroupsResponse(data=await self.get_all_with_type("tool_group"))
+
+    async def get_tool_group(self, toolgroup_id: str) -> ToolGroup:
+        tool_group = await self.get_object_by_identifier("tool_group", toolgroup_id)
+        if tool_group is None:
+            raise ValueError(f"Tool group '{toolgroup_id}' not found")
+        return tool_group
+
+    async def get_tool(self, tool_name: str) -> Tool:
+        return await self.get_object_by_identifier("tool", tool_name)
+
+    async def register_tool_group(
+        self,
+        toolgroup_id: str,
+        provider_id: str,
+        mcp_endpoint: URL | None = None,
+        args: dict[str, Any] | None = None,
+    ) -> None:
+        tools = []
+        tool_defs = await self.impls_by_provider_id[provider_id].list_runtime_tools(toolgroup_id, mcp_endpoint)
+        tool_host = ToolHost.model_context_protocol if mcp_endpoint else ToolHost.distribution
+
+        for tool_def in tool_defs.data:
+            tools.append(
+                ToolWithACL(
+                    identifier=tool_def.name,
+                    toolgroup_id=toolgroup_id,
+                    description=tool_def.description or "",
+                    parameters=tool_def.parameters or [],
+                    provider_id=provider_id,
+                    provider_resource_id=tool_def.name,
+                    metadata=tool_def.metadata,
+                    tool_host=tool_host,
+                )
+            )
+        for tool in tools:
+            existing_tool = await self.get_tool(tool.identifier)
+            # Compare existing and new object if one exists
+            if existing_tool:
+                existing_dict = existing_tool.model_dump()
+                new_dict = tool.model_dump()
+
+                if existing_dict != new_dict:
+                    raise ValueError(
+                        f"Object {tool.identifier} already exists in registry. Please use a different identifier."
+                    )
+            await self.register_object(tool)
+
+        await self.dist_registry.register(
+            ToolGroupWithACL(
+                identifier=toolgroup_id,
+                provider_id=provider_id,
+                provider_resource_id=toolgroup_id,
+                mcp_endpoint=mcp_endpoint,
+                args=args,
+            )
+        )
+
+    async def unregister_toolgroup(self, toolgroup_id: str) -> None:
+        tool_group = await self.get_tool_group(toolgroup_id)
+        if tool_group is None:
+            raise ValueError(f"Tool group {toolgroup_id} not found")
+        tools = await self.list_tools(toolgroup_id)
+        for tool in getattr(tools, "data", []):
+            await self.unregister_object(tool)
+        await self.unregister_object(tool_group)
+
+    async def shutdown(self) -> None:
+        pass