slight rename

This commit is contained in:
Ashwin Bharambe 2025-01-21 15:49:22 -08:00
parent 68f2550e1e
commit 039982004d
2 changed files with 4 additions and 4 deletions

View file

@ -132,7 +132,7 @@ class ToolGroups(Protocol):
... ...
class SpecialToolGroups(Enum): class SpecialToolGroup(Enum):
rag_tool = "rag_tool" rag_tool = "rag_tool"

View file

@ -9,7 +9,7 @@ from typing import Dict, List
from pydantic import BaseModel from pydantic import BaseModel
from llama_stack.apis.tools import RAGToolRuntime, SpecialToolGroups from llama_stack.apis.tools import RAGToolRuntime, SpecialToolGroup
from llama_stack.apis.version import LLAMA_STACK_API_VERSION from llama_stack.apis.version import LLAMA_STACK_API_VERSION
@ -26,7 +26,7 @@ class ApiEndpoint(BaseModel):
def toolgroup_protocol_map(): def toolgroup_protocol_map():
return { return {
SpecialToolGroups.rag_tool: RAGToolRuntime, SpecialToolGroup.rag_tool: RAGToolRuntime,
} }
@ -39,7 +39,7 @@ def get_all_api_endpoints() -> Dict[Api, List[ApiEndpoint]]:
endpoints = [] endpoints = []
protocol_methods = inspect.getmembers(protocol, predicate=inspect.isfunction) protocol_methods = inspect.getmembers(protocol, predicate=inspect.isfunction)
if api == Api.tool_runtime: if api == Api.tool_runtime:
for tool_group in SpecialToolGroups: for tool_group in SpecialToolGroup:
sub_protocol = toolgroup_protocols[tool_group] sub_protocol = toolgroup_protocols[tool_group]
sub_protocol_methods = inspect.getmembers( sub_protocol_methods = inspect.getmembers(
sub_protocol, predicate=inspect.isfunction sub_protocol, predicate=inspect.isfunction