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

fix: non-llama tool extraction

Browse source 
Summary:

Test Plan:
Summary:

Test Plan:
Summary:

Test Plan:
This commit is contained in:
Eric Huang 2025-03-17 15:08:40 -07:00
parent 60ae7455f6
commit 71db8cd074
4 changed files with 161 additions and 125 deletions
Download patch file Download diff file Expand all files Collapse all files

73
tests/integration/agents/test_agents.py
View file

@ -271,7 +271,7 @@ def test_custom_tool(llama_stack_client_with_mocked_inference, agent_config):
client_tool = get_boiling_point
agent_config = {
**agent_config,
"tools": ["builtin::websearch", client_tool],
"tools": [client_tool],
}
agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
@ -320,42 +320,55 @@ def test_custom_tool_infinite_loop(llama_stack_client_with_mocked_inference, age
assert num_tool_calls <= 5
def test_tool_choice(llama_stack_client_with_mocked_inference, agent_config):
def run_agent(tool_choice):
client_tool = get_boiling_point
test_agent_config = {
**agent_config,
"tool_config": {"tool_choice": tool_choice},
"tools": [client_tool],
}
agent = Agent(llama_stack_client_with_mocked_inference, **test_agent_config)
session_id = agent.create_session(f"test-session-{uuid4()}")
response = agent.create_turn(
messages=[
{
"role": "user",
"content": "What is the boiling point of polyjuice?",
},
],
session_id=session_id,
stream=False,
)
return [step for step in response.steps if step.step_type == "tool_execution"]
tool_execution_steps = run_agent("required")
def test_tool_choice_required(llama_stack_client_with_mocked_inference, agent_config):
tool_execution_steps = run_agent_with_tool_choice(
llama_stack_client_with_mocked_inference, agent_config, "required"
)
assert len(tool_execution_steps) > 0
tool_execution_steps = run_agent("none")
def test_tool_choice_none(llama_stack_client_with_mocked_inference, agent_config):
tool_execution_steps = run_agent_with_tool_choice(llama_stack_client_with_mocked_inference, agent_config, "none")
assert len(tool_execution_steps) == 0
tool_execution_steps = run_agent("get_boiling_point")
def test_tool_choice_get_boiling_point(llama_stack_client_with_mocked_inference, agent_config):
if "llama" not in agent_config["model"].lower():
pytest.xfail("NotImplemented for non-llama models")
tool_execution_steps = run_agent_with_tool_choice(
llama_stack_client_with_mocked_inference, agent_config, "get_boiling_point"
)
assert len(tool_execution_steps) >= 1 and tool_execution_steps[0].tool_calls[0].tool_name == "get_boiling_point"
def run_agent_with_tool_choice(client, agent_config, tool_choice):
client_tool = get_boiling_point
test_agent_config = {
**agent_config,
"tool_config": {"tool_choice": tool_choice},
"tools": [client_tool],
"max_infer_iters": 2,
}
agent = Agent(client, **test_agent_config)
session_id = agent.create_session(f"test-session-{uuid4()}")
response = agent.create_turn(
messages=[
{
"role": "user",
"content": "What is the boiling point of polyjuice?",
},
],
session_id=session_id,
stream=False,
)
return [step for step in response.steps if step.step_type == "tool_execution"]
@pytest.mark.parametrize("rag_tool_name", ["builtin::rag/knowledge_search", "builtin::rag"])
def test_rag_agent(llama_stack_client_with_mocked_inference, agent_config, rag_tool_name):
urls = ["chat.rst", "llama3.rst", "memory_optimizations.rst", "lora_finetune.rst"]