mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-28 02:53:30 +00:00
chore: refactor client tool in test (#1331)
# What does this PR do? Use @client_tool decorator instead of ClientTool [//]: # (If resolving an issue, uncomment and update the line below) [//]: # (Closes #[issue-number]) ## Test Plan ``` LLAMA_STACK_CONFIG=fireworks pytest -v tests/client-sdk/agents/test_agents.py --inference-model "meta-llama/Llama-3.3-70B-Instruct" ``` <img width="1053" alt="image" src="https://github.com/user-attachments/assets/d3ade884-ef42-494e-8028-3b09d9ef1978" /> [//]: # (## Documentation)
This commit is contained in:
parent
15f69e75ff
commit
82fa0803fa
1 changed files with 18 additions and 64 deletions
|
@ -4,20 +4,15 @@
|
|||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
import json
|
||||
from typing import Dict, List
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from llama_stack_client.lib.agents.agent import Agent
|
||||
from llama_stack_client.lib.agents.client_tool import ClientTool
|
||||
from llama_stack_client.lib.agents.client_tool import client_tool
|
||||
from llama_stack_client.lib.agents.event_logger import EventLogger
|
||||
from llama_stack_client.types import ToolResponseMessage
|
||||
from llama_stack_client.types.agents.turn_create_params import Document as AgentDocument
|
||||
from llama_stack_client.types.memory_insert_params import Document
|
||||
from llama_stack_client.types.shared.completion_message import CompletionMessage
|
||||
from llama_stack_client.types.shared_params.agent_config import AgentConfig, ToolConfig
|
||||
from llama_stack_client.types.tool_def_param import Parameter
|
||||
|
||||
from llama_stack.apis.agents.agents import (
|
||||
AgentConfig as Server__AgentConfig,
|
||||
|
@ -27,56 +22,15 @@ from llama_stack.apis.agents.agents import (
|
|||
)
|
||||
|
||||
|
||||
class TestClientTool(ClientTool):
|
||||
"""Tool to give boiling point of a liquid
|
||||
Returns the correct value for polyjuice in Celcius and Fahrenheit
|
||||
and returns -1 for other liquids
|
||||
@client_tool
|
||||
def get_boiling_point(liquid_name: str, celcius: bool = True) -> int:
|
||||
"""
|
||||
Returns the boiling point of a liquid in Celcius or Fahrenheit
|
||||
|
||||
def run(self, messages: List[CompletionMessage]) -> List[ToolResponseMessage]:
|
||||
assert len(messages) == 1, "Expected single message"
|
||||
|
||||
message = messages[0]
|
||||
|
||||
tool_call = message.tool_calls[0]
|
||||
|
||||
try:
|
||||
response = self.run_impl(**tool_call.arguments)
|
||||
response_str = json.dumps(response, ensure_ascii=False)
|
||||
except Exception as e:
|
||||
response_str = f"Error when running tool: {e}"
|
||||
|
||||
message = ToolResponseMessage(
|
||||
role="tool",
|
||||
call_id=tool_call.call_id,
|
||||
tool_name=tool_call.tool_name,
|
||||
content=response_str,
|
||||
)
|
||||
return message
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "get_boiling_point"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "Get the boiling point of imaginary liquids (eg. polyjuice)"
|
||||
|
||||
def get_params_definition(self) -> Dict[str, Parameter]:
|
||||
return {
|
||||
"liquid_name": Parameter(
|
||||
name="liquid_name",
|
||||
parameter_type="string",
|
||||
description="The name of the liquid",
|
||||
required=True,
|
||||
),
|
||||
"celcius": Parameter(
|
||||
name="celcius",
|
||||
parameter_type="boolean",
|
||||
description="Whether to return the boiling point in Celcius",
|
||||
required=False,
|
||||
),
|
||||
}
|
||||
|
||||
def run_impl(self, liquid_name: str, celcius: bool = True) -> int:
|
||||
:param liquid_name: The name of the liquid
|
||||
:param celcius: Whether to return the boiling point in Celcius
|
||||
:return: The boiling point of the liquid in Celcius or Fahrenheit
|
||||
"""
|
||||
if liquid_name.lower() == "polyjuice":
|
||||
if celcius:
|
||||
return -100
|
||||
|
@ -298,7 +252,7 @@ def test_code_interpreter_for_attachments(llama_stack_client, agent_config):
|
|||
|
||||
|
||||
def test_custom_tool(llama_stack_client, agent_config):
|
||||
client_tool = TestClientTool()
|
||||
client_tool = get_boiling_point
|
||||
agent_config = {
|
||||
**agent_config,
|
||||
"toolgroups": ["builtin::websearch"],
|
||||
|
@ -326,7 +280,7 @@ def test_custom_tool(llama_stack_client, agent_config):
|
|||
|
||||
def test_tool_choice(llama_stack_client, agent_config):
|
||||
def run_agent(tool_choice):
|
||||
client_tool = TestClientTool()
|
||||
client_tool = get_boiling_point
|
||||
|
||||
test_agent_config = {
|
||||
**agent_config,
|
||||
|
@ -362,7 +316,7 @@ def test_tool_choice(llama_stack_client, agent_config):
|
|||
|
||||
# TODO: fix this flaky test
|
||||
def xtest_override_system_message_behavior(llama_stack_client, agent_config):
|
||||
client_tool = TestClientTool()
|
||||
client_tool = get_boiling_point
|
||||
agent_config = {
|
||||
**agent_config,
|
||||
"instructions": "You are a pirate",
|
||||
|
@ -586,7 +540,7 @@ def test_rag_and_code_agent(llama_stack_client, agent_config):
|
|||
|
||||
|
||||
def test_create_turn_response(llama_stack_client, agent_config):
|
||||
client_tool = TestClientTool()
|
||||
client_tool = get_boiling_point
|
||||
agent_config = {
|
||||
**agent_config,
|
||||
"input_shields": [],
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue