forked from phoenix-oss/llama-stack-mirror
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
|
# This source code is licensed under the terms described in the LICENSE file in
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
import json
|
|
||||||
from typing import Dict, List
|
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from llama_stack_client.lib.agents.agent import Agent
|
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.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.agents.turn_create_params import Document as AgentDocument
|
||||||
from llama_stack_client.types.memory_insert_params import Document
|
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.shared_params.agent_config import AgentConfig, ToolConfig
|
||||||
from llama_stack_client.types.tool_def_param import Parameter
|
|
||||||
|
|
||||||
from llama_stack.apis.agents.agents import (
|
from llama_stack.apis.agents.agents import (
|
||||||
AgentConfig as Server__AgentConfig,
|
AgentConfig as Server__AgentConfig,
|
||||||
|
@ -27,63 +22,22 @@ from llama_stack.apis.agents.agents import (
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class TestClientTool(ClientTool):
|
@client_tool
|
||||||
"""Tool to give boiling point of a liquid
|
def get_boiling_point(liquid_name: str, celcius: bool = True) -> int:
|
||||||
Returns the correct value for polyjuice in Celcius and Fahrenheit
|
|
||||||
and returns -1 for other liquids
|
|
||||||
"""
|
"""
|
||||||
|
Returns the boiling point of a liquid in Celcius or Fahrenheit
|
||||||
|
|
||||||
def run(self, messages: List[CompletionMessage]) -> List[ToolResponseMessage]:
|
:param liquid_name: The name of the liquid
|
||||||
assert len(messages) == 1, "Expected single message"
|
:param celcius: Whether to return the boiling point in Celcius
|
||||||
|
:return: The boiling point of the liquid in Celcius or Fahrenheit
|
||||||
message = messages[0]
|
"""
|
||||||
|
if liquid_name.lower() == "polyjuice":
|
||||||
tool_call = message.tool_calls[0]
|
if celcius:
|
||||||
|
return -100
|
||||||
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:
|
|
||||||
if liquid_name.lower() == "polyjuice":
|
|
||||||
if celcius:
|
|
||||||
return -100
|
|
||||||
else:
|
|
||||||
return -212
|
|
||||||
else:
|
else:
|
||||||
return -1
|
return -212
|
||||||
|
else:
|
||||||
|
return -1
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="session")
|
@pytest.fixture(scope="session")
|
||||||
|
@ -298,7 +252,7 @@ def test_code_interpreter_for_attachments(llama_stack_client, agent_config):
|
||||||
|
|
||||||
|
|
||||||
def test_custom_tool(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 = {
|
||||||
**agent_config,
|
**agent_config,
|
||||||
"toolgroups": ["builtin::websearch"],
|
"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 test_tool_choice(llama_stack_client, agent_config):
|
||||||
def run_agent(tool_choice):
|
def run_agent(tool_choice):
|
||||||
client_tool = TestClientTool()
|
client_tool = get_boiling_point
|
||||||
|
|
||||||
test_agent_config = {
|
test_agent_config = {
|
||||||
**agent_config,
|
**agent_config,
|
||||||
|
@ -362,7 +316,7 @@ def test_tool_choice(llama_stack_client, agent_config):
|
||||||
|
|
||||||
# TODO: fix this flaky test
|
# TODO: fix this flaky test
|
||||||
def xtest_override_system_message_behavior(llama_stack_client, agent_config):
|
def xtest_override_system_message_behavior(llama_stack_client, agent_config):
|
||||||
client_tool = TestClientTool()
|
client_tool = get_boiling_point
|
||||||
agent_config = {
|
agent_config = {
|
||||||
**agent_config,
|
**agent_config,
|
||||||
"instructions": "You are a pirate",
|
"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):
|
def test_create_turn_response(llama_stack_client, agent_config):
|
||||||
client_tool = TestClientTool()
|
client_tool = get_boiling_point
|
||||||
agent_config = {
|
agent_config = {
|
||||||
**agent_config,
|
**agent_config,
|
||||||
"input_shields": [],
|
"input_shields": [],
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue