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

add support for built in tool type

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-12-23 16:50:03 -08:00
parent 517bc9ebea
commit 1a66ddc1b5
8 changed files with 83 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -621,6 +621,9 @@ class ChatAgent(ShieldRunnerMixin):
ret = []
for tool_name in self.agent_config.available_tools:
tool = await self.tool_groups_api.get_tool(tool_name)
if tool.built_in_type:
ret.append(ToolDefinition(tool_name=tool.built_in_type))
continue
params = {}
for param in tool.parameters:
params[param.name] = ToolParamDefinition(

3
llama_stack/providers/inline/tool_runtime/brave_search/brave_search.py
View file

@ -25,8 +25,7 @@ class BraveSearchToolRuntimeImpl(
pass
async def register_tool(self, tool: Tool):
if tool.identifier != "brave_search":
raise ValueError(f"Tool identifier {tool.identifier} is not supported")
pass
async def unregister_tool(self, tool_id: str) -> None:
return

3
llama_stack/providers/inline/tool_runtime/tavily_search/tavily_search.py
View file

@ -26,8 +26,7 @@ class TavilySearchToolRuntimeImpl(
pass
async def register_tool(self, tool: Tool):
if tool.identifier != "tavily_search":
raise ValueError(f"Tool identifier {tool.identifier} is not supported")
pass
async def unregister_tool(self, tool_id: str) -> None:
return

3
llama_stack/providers/remote/tool_runtime/model_context_protocol/model_context_protocol.py
View file

@ -8,6 +8,7 @@ from typing import Any, Dict, List
from urllib.parse import urlparse
from llama_stack.apis.tools import (
CustomToolDef,
MCPToolGroupDef,
ToolDef,
ToolGroupDef,
@ -52,7 +53,7 @@ class ModelContextProtocolToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
)
)
tools.append(
ToolDef(
CustomToolDef(
name=tool.name,
description=tool.description,
parameters=parameters,

46
llama_stack/providers/tests/agents/fixtures.py
View file

@ -9,10 +9,12 @@ import tempfile
import pytest
import pytest_asyncio
from llama_models.llama3.api.datatypes import BuiltinTool
from llama_stack.apis.models import ModelInput, ModelType
from llama_stack.apis.tools import (
ToolDef,
BuiltInToolDef,
CustomToolDef,
ToolGroupInput,
ToolParameter,
UserDefinedToolGroupDef,
@ -151,42 +153,12 @@ async def agents_stack(request, inference_model, safety_shield):
)
)
tool_groups = [
ToolGroupInput(
tool_group_id="brave_search_group",
tool_group=UserDefinedToolGroupDef(
tools=[
ToolDef(
name="brave_search",
description="brave_search",
parameters=[
ToolParameter(
name="query",
description="query",
parameter_type="string",
required=True,
),
],
metadata={},
),
],
),
provider_id="brave-search",
),
ToolGroupInput(
tool_group_id="tavily_search_group",
tool_group=UserDefinedToolGroupDef(
tools=[
ToolDef(
name="tavily_search",
description="tavily_search",
parameters=[
ToolParameter(
name="query",
description="query",
parameter_type="string",
required=True,
),
],
BuiltInToolDef(
built_in_type=BuiltinTool.brave_search,
metadata={},
),
],
@ -197,7 +169,7 @@ async def agents_stack(request, inference_model, safety_shield):
tool_group_id="memory_group",
tool_group=UserDefinedToolGroupDef(
tools=[
ToolDef(
CustomToolDef(
name="memory",
description="memory",
parameters=[
@ -230,10 +202,8 @@ async def agents_stack(request, inference_model, safety_shield):
tool_group_id="code_interpreter_group",
tool_group=UserDefinedToolGroupDef(
tools=[
ToolDef(
name="code_interpreter",
description="code_interpreter",
parameters=[],
BuiltInToolDef(
built_in_type=BuiltinTool.code_interpreter,
metadata={},
)
],

18
llama_stack/providers/tests/agents/test_agents.py
View file

@ -150,9 +150,7 @@ async def create_agent_turn_with_search_tool(
assert isinstance(tool_execution, ToolExecutionStep)
assert len(tool_execution.tool_calls) > 0
actual_tool_name = tool_execution.tool_calls[0].tool_name
if isinstance(actual_tool_name, BuiltinTool):
actual_tool_name = actual_tool_name.value
assert actual_tool_name == tool_name
assert actual_tool_name.value == tool_name
assert len(tool_execution.tool_responses) > 0
check_turn_complete_event(turn_response, session_id, search_query_messages)
@ -305,20 +303,6 @@ class TestAgents:
"brave_search",
)
@pytest.mark.asyncio
async def test_create_agent_turn_with_tavily_search(
self, agents_stack, search_query_messages, common_params
):
if "TAVILY_SEARCH_API_KEY" not in os.environ:
pytest.skip("TAVILY_SEARCH_API_KEY not set, skipping test")
await create_agent_turn_with_search_tool(
agents_stack,
search_query_messages,
common_params,
"tavily_search",
)
def check_event_types(turn_response):
event_types = [chunk.event.payload.event_type for chunk in turn_response]