phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-09 13:14:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(tools): use { input_schema, output_schema } for ToolDefinition

Browse source
This commit is contained in:
Ashwin Bharambe 2025-09-30 19:13:15 -07:00
parent 42414a1a1b
commit 139320e19f
20 changed files with 1989 additions and 386 deletions
Download patch file Download diff file Expand all files Collapse all files

369
tests/integration/inference/test_tools_with_schemas.py Normal file
View file

@ -0,0 +1,369 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
"""
Integration tests for inference/chat completion with JSON Schema-based tools.
Tests that tools pass through correctly to various LLM providers.
"""
import json
import pytest
from llama_stack import LlamaStackAsLibraryClient
from llama_stack.models.llama.datatypes import ToolDefinition
from tests.common.mcp import make_mcp_server
AUTH_TOKEN = "test-token"
class TestChatCompletionWithTools:
"""Test chat completion with tools that have complex schemas."""
def test_simple_tool_call(self, llama_stack_client, text_model_id):
"""Test basic tool calling with simple input schema."""
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string", "description": "City name"}},
"required": ["location"],
},
},
}
]
response = llama_stack_client.chat.completions.create(
model=text_model_id,
messages=[{"role": "user", "content": "What's the weather in San Francisco?"}],
tools=tools,
)
assert response is not None
def test_tool_with_complex_schema(self, llama_stack_client, text_model_id):
"""Test tool calling with complex schema including $ref and $defs."""
tools = [
{
"type": "function",
"function": {
"name": "book_flight",
"description": "Book a flight",
"parameters": {
"type": "object",
"properties": {
"flight": {"$ref": "#/$defs/FlightInfo"},
"passenger": {"$ref": "#/$defs/Passenger"},
},
"required": ["flight", "passenger"],
"$defs": {
"FlightInfo": {
"type": "object",
"properties": {
"from": {"type": "string"},
"to": {"type": "string"},
"date": {"type": "string", "format": "date"},
},
},
"Passenger": {
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
},
},
},
},
}
]
response = llama_stack_client.chat.completions.create(
model=text_model_id,
messages=[{"role": "user", "content": "Book a flight from SFO to JFK for John Doe"}],
tools=tools,
)
# The key test: No errors during schema processing
# The LLM received a valid, complete schema with $ref/$defs
assert response is not None
class TestOpenAICompatibility:
"""Test OpenAI-compatible endpoints with new schema format."""
def test_openai_chat_completion_with_tools(self, compat_client, text_model_id):
"""Test OpenAI-compatible chat completion with tools."""
from openai import OpenAI
if not isinstance(compat_client, OpenAI):
pytest.skip("OpenAI client required")
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string", "description": "City name"}},
"required": ["location"],
},
},
}
]
response = compat_client.chat.completions.create(
model=text_model_id, messages=[{"role": "user", "content": "What's the weather in Tokyo?"}], tools=tools
)
assert response is not None
assert response.choices is not None
def test_openai_format_preserves_complex_schemas(self, compat_client, text_model_id):
"""Test that complex schemas work through OpenAI-compatible API."""
from openai import OpenAI
if not isinstance(compat_client, OpenAI):
pytest.skip("OpenAI client required")
tools = [
{
"type": "function",
"function": {
"name": "process_data",
"description": "Process structured data",
"parameters": {
"type": "object",
"properties": {"data": {"$ref": "#/$defs/DataObject"}},
"$defs": {
"DataObject": {
"type": "object",
"properties": {"values": {"type": "array", "items": {"type": "number"}}},
}
},
},
},
}
]
response = compat_client.chat.completions.create(
model=text_model_id, messages=[{"role": "user", "content": "Process this data"}], tools=tools
)
assert response is not None
class TestMCPToolsInChatCompletion:
"""Test using MCP tools in chat completion."""
@pytest.fixture
def mcp_with_schemas(self):
"""MCP server for chat completion tests."""
from mcp.server.fastmcp import Context
async def calculate(x: float, y: float, operation: str, ctx: Context) -> float:
ops = {"add": x + y, "sub": x - y, "mul": x * y, "div": x / y if y != 0 else None}
return ops.get(operation, 0)
with make_mcp_server(required_auth_token=AUTH_TOKEN, tools={"calculate": calculate}) as server:
yield server
def test_mcp_tools_in_inference(self, llama_stack_client, text_model_id, mcp_with_schemas):
"""Test that MCP tools can be used in inference."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::calc"
uri = mcp_with_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# Get the tools from MCP
tools_response = llama_stack_client.tool_runtime.list_runtime_tools(
tool_group_id=test_toolgroup_id,
extra_headers=auth_headers,
)
# Convert to OpenAI format for inference
tools = []
for tool in tools_response.data:
tools.append(
{
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": tool.input_schema if hasattr(tool, "input_schema") else {},
},
}
)
# Use in chat completion
response = llama_stack_client.chat.completions.create(
model=text_model_id,
messages=[{"role": "user", "content": "Calculate 5 + 3"}],
tools=tools,
)
# Schema should have been passed through correctly
assert response is not None
class TestProviderSpecificBehavior:
"""Test provider-specific handling of schemas."""
def test_openai_provider_drops_output_schema(self, llama_stack_client, text_model_id):
"""Test that OpenAI provider doesn't send output_schema (API limitation)."""
# This is more of a documentation test
# OpenAI API doesn't support output schemas, so we drop them
_tool = ToolDefinition(
tool_name="test",
input_schema={"type": "object", "properties": {"x": {"type": "string"}}},
output_schema={"type": "object", "properties": {"y": {"type": "number"}}},
)
# When this tool is sent to OpenAI provider, output_schema is dropped
# But input_schema is preserved
# This test documents the expected behavior
# We can't easily test this without mocking, but the unit tests cover it
pass
def test_gemini_array_support(self):
"""Test that Gemini receives array schemas correctly (issue from commit 65f7b81e)."""
# This was the original bug that led to adding 'items' field
# Now with full JSON Schema pass-through, arrays should work
tool = ToolDefinition(
tool_name="tag_processor",
input_schema={
"type": "object",
"properties": {"tags": {"type": "array", "items": {"type": "string"}, "description": "List of tags"}},
},
)
# With new approach, the complete schema with items is preserved
assert tool.input_schema["properties"]["tags"]["type"] == "array"
assert tool.input_schema["properties"]["tags"]["items"]["type"] == "string"
class TestStreamingWithTools:
"""Test streaming chat completion with tools."""
def test_streaming_tool_calls(self, llama_stack_client, text_model_id):
"""Test that tool schemas work correctly in streaming mode."""
tools = [
{
"type": "function",
"function": {
"name": "get_time",
"description": "Get current time",
"parameters": {"type": "object", "properties": {"timezone": {"type": "string"}}},
},
}
]
response_stream = llama_stack_client.chat.completions.create(
model=text_model_id,
messages=[{"role": "user", "content": "What time is it in UTC?"}],
tools=tools,
stream=True,
)
# Should be able to iterate through stream
chunks = []
for chunk in response_stream:
chunks.append(chunk)
# Should have received some chunks
assert len(chunks) >= 0
class TestEdgeCases:
"""Test edge cases in inference with tools."""
def test_tool_without_schema(self, llama_stack_client, text_model_id):
"""Test tool with no input_schema."""
tools = [
{
"type": "function",
"function": {
"name": "no_args_tool",
"description": "Tool with no arguments",
"parameters": {"type": "object", "properties": {}},
},
}
]
response = llama_stack_client.chat.completions.create(
model=text_model_id,
messages=[{"role": "user", "content": "Call the no args tool"}],
tools=tools,
)
assert response is not None
def test_multiple_tools_with_different_schemas(self, llama_stack_client, text_model_id):
"""Test multiple tools with different schema complexities."""
tools = [
{
"type": "function",
"function": {
"name": "simple",
"parameters": {"type": "object", "properties": {"x": {"type": "string"}}},
},
},
{
"type": "function",
"function": {
"name": "complex",
"parameters": {
"type": "object",
"properties": {"data": {"$ref": "#/$defs/Complex"}},
"$defs": {
"Complex": {
"type": "object",
"properties": {"nested": {"type": "array", "items": {"type": "number"}}},
}
},
},
},
},
{
"type": "function",
"function": {
"name": "with_output",
"parameters": {"type": "object", "properties": {"input": {"type": "string"}}},
},
},
]
response = llama_stack_client.chat.completions.create(
model=text_model_id,
messages=[{"role": "user", "content": "Use one of the available tools"}],
tools=tools,
)
# All tools should have been processed without errors
assert response is not None

478
tests/integration/tool_runtime/test_mcp_json_schema.py Normal file
View file

@ -0,0 +1,478 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
"""
Integration tests for MCP tools with complex JSON Schema support.
Tests $ref, $defs, and other JSON Schema features through MCP integration.
"""
import json
import pytest
from llama_stack import LlamaStackAsLibraryClient
from tests.common.mcp import make_mcp_server
AUTH_TOKEN = "test-token"
@pytest.fixture(scope="function")
def mcp_server_with_complex_schemas():
"""MCP server with tools that have complex schemas including $ref and $defs."""
from mcp.server.fastmcp import Context
async def book_flight(flight: dict, passengers: list[dict], payment: dict, ctx: Context) -> dict:
"""
Book a flight with passenger and payment information.
This tool uses JSON Schema $ref and $defs for type reuse.
"""
return {
"booking_id": "BK12345",
"flight": flight,
"passengers": passengers,
"payment": payment,
"status": "confirmed",
}
async def process_order(order_data: dict, ctx: Context) -> dict:
"""
Process an order with nested address information.
Uses nested objects and $ref.
"""
return {"order_id": "ORD789", "status": "processing", "data": order_data}
async def flexible_contact(contact_info: str, ctx: Context) -> dict:
"""
Accept flexible contact (email or phone).
Uses anyOf schema.
"""
if "@" in contact_info:
return {"type": "email", "value": contact_info}
else:
return {"type": "phone", "value": contact_info}
# Manually attach complex schemas to the functions
# (FastMCP might not support this by default, so this is test setup)
# For MCP, we need to set the schema via tool annotations
# This is test infrastructure to force specific schemas
tools = {"book_flight": book_flight, "process_order": process_order, "flexible_contact": flexible_contact}
# Note: In real MCP implementation, we'd configure these schemas properly
# For testing, we may need to mock or extend the MCP server setup
with make_mcp_server(required_auth_token=AUTH_TOKEN, tools=tools) as server_info:
yield server_info
@pytest.fixture(scope="function")
def mcp_server_with_output_schemas():
"""MCP server with tools that have output schemas defined."""
from mcp.server.fastmcp import Context
async def get_weather(location: str, ctx: Context) -> dict:
"""
Get weather with structured output.
Has both input and output schemas.
"""
return {"temperature": 72.5, "conditions": "Sunny", "humidity": 45, "wind_speed": 10.2}
async def calculate(x: float, y: float, operation: str, ctx: Context) -> dict:
"""
Perform calculation with validated output.
"""
operations = {"add": x + y, "subtract": x - y, "multiply": x * y, "divide": x / y if y != 0 else None}
result = operations.get(operation)
return {"result": result, "operation": operation}
tools = {"get_weather": get_weather, "calculate": calculate}
with make_mcp_server(required_auth_token=AUTH_TOKEN, tools=tools) as server_info:
yield server_info
class TestMCPSchemaPreservation:
"""Test that MCP tool schemas are preserved correctly."""
def test_mcp_tools_list_with_schemas(self, llama_stack_client, mcp_server_with_complex_schemas):
"""Test listing MCP tools preserves input_schema."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::complex"
uri = mcp_server_with_complex_schemas["server_url"]
# Clean up any existing registration
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
# Register MCP toolgroup
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# List runtime tools
response = llama_stack_client.tool_runtime.list_runtime_tools(
tool_group_id=test_toolgroup_id,
extra_headers=auth_headers,
)
tools = response.data
assert len(tools) > 0
# Check each tool has input_schema
for tool in tools:
assert hasattr(tool, "input_schema")
# Schema might be None or a dict depending on tool
if tool.input_schema is not None:
assert isinstance(tool.input_schema, dict)
# Should have basic JSON Schema structure
if "properties" in tool.input_schema:
assert "type" in tool.input_schema
def test_mcp_schema_with_refs_preserved(self, llama_stack_client, mcp_server_with_complex_schemas):
"""Test that $ref and $defs in MCP schemas are preserved."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::complex"
uri = mcp_server_with_complex_schemas["server_url"]
# Register
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# List tools
response = llama_stack_client.tool_runtime.list_runtime_tools(
tool_group_id=test_toolgroup_id,
extra_headers=auth_headers,
)
# Find book_flight tool (which should have $ref/$defs)
book_flight_tool = next((t for t in response.data if t.name == "book_flight"), None)
if book_flight_tool and book_flight_tool.input_schema:
# If the MCP server provides $defs, they should be preserved
# This is the KEY test for the bug fix
schema = book_flight_tool.input_schema
# Check if schema has properties (might vary based on MCP implementation)
if "properties" in schema:
# Verify schema structure is preserved (exact structure depends on MCP server)
assert isinstance(schema["properties"], dict)
# If $defs are present, verify they're preserved
if "$defs" in schema:
assert isinstance(schema["$defs"], dict)
# Each definition should be a dict
for _def_name, def_schema in schema["$defs"].items():
assert isinstance(def_schema, dict)
def test_mcp_output_schema_preserved(self, llama_stack_client, mcp_server_with_output_schemas):
"""Test that MCP outputSchema is preserved."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::with_output"
uri = mcp_server_with_output_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
response = llama_stack_client.tool_runtime.list_runtime_tools(
tool_group_id=test_toolgroup_id,
extra_headers=auth_headers,
)
# Find get_weather tool
weather_tool = next((t for t in response.data if t.name == "get_weather"), None)
if weather_tool:
# Check if output_schema field exists and is preserved
assert hasattr(weather_tool, "output_schema")
# If MCP server provides output schema, it should be preserved
if weather_tool.output_schema is not None:
assert isinstance(weather_tool.output_schema, dict)
# Should have JSON Schema structure
if "properties" in weather_tool.output_schema:
assert "type" in weather_tool.output_schema
class TestMCPToolInvocation:
"""Test invoking MCP tools with complex schemas."""
def test_invoke_mcp_tool_with_nested_data(self, llama_stack_client, mcp_server_with_complex_schemas):
"""Test invoking MCP tool that expects nested object structure."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::complex"
uri = mcp_server_with_complex_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# Invoke tool with complex nested data
result = llama_stack_client.tool_runtime.invoke_tool(
tool_name="process_order",
kwargs={
"order_data": {
"items": [{"name": "Widget", "quantity": 2}, {"name": "Gadget", "quantity": 1}],
"shipping": {"address": {"street": "123 Main St", "city": "San Francisco", "zipcode": "94102"}},
}
},
extra_headers=auth_headers,
)
# Should succeed without schema validation errors
assert result.content is not None
assert result.error_message is None
def test_invoke_with_flexible_schema(self, llama_stack_client, mcp_server_with_complex_schemas):
"""Test invoking tool with anyOf schema (flexible input)."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::complex"
uri = mcp_server_with_complex_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# Test with email format
result_email = llama_stack_client.tool_runtime.invoke_tool(
tool_name="flexible_contact",
kwargs={"contact_info": "user@example.com"},
extra_headers=auth_headers,
)
assert result_email.error_message is None
# Test with phone format
result_phone = llama_stack_client.tool_runtime.invoke_tool(
tool_name="flexible_contact",
kwargs={"contact_info": "+15551234567"},
extra_headers=auth_headers,
)
assert result_phone.error_message is None
class TestAgentWithMCPTools:
"""Test agents using MCP tools with complex schemas."""
def test_agent_with_complex_mcp_tool(self, llama_stack_client, text_model_id, mcp_server_with_complex_schemas):
"""Test agent can use MCP tools with $ref/$defs schemas."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
from llama_stack_client import Agent
test_toolgroup_id = "mcp::complex"
uri = mcp_server_with_complex_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# Create agent with MCP tools
agent = Agent(
client=llama_stack_client,
model=text_model_id,
instructions="You are a helpful assistant that can process orders and book flights.",
tools=[test_toolgroup_id],
)
session_id = agent.create_session("test-session-complex")
# Ask agent to use a tool with complex schema
response = agent.create_turn(
session_id=session_id,
messages=[
{"role": "user", "content": "Process an order with 2 widgets going to 123 Main St, San Francisco"}
],
stream=False,
extra_headers=auth_headers,
)
steps = response.steps
# Verify agent was able to call the tool
# (The LLM should have been able to understand the schema and formulate a valid call)
tool_execution_steps = [s for s in steps if s.step_type == "tool_execution"]
# Agent might or might not call the tool depending on the model
# But if it does, there should be no errors
for step in tool_execution_steps:
if step.tool_responses:
for tool_response in step.tool_responses:
assert tool_response.content is not None
class TestSchemaValidation:
"""Test schema validation (future feature)."""
def test_invalid_input_rejected(self, llama_stack_client, mcp_server_with_complex_schemas):
"""Test that invalid input is rejected (if validation is implemented)."""
# This test documents expected behavior once we add input validation
# For now, it may pass invalid data through
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::complex"
uri = mcp_server_with_complex_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# Try to invoke with completely wrong data type
# Once validation is added, this should raise an error
try:
llama_stack_client.tool_runtime.invoke_tool(
tool_name="process_order",
kwargs={"order_data": "this should be an object not a string"},
extra_headers=auth_headers,
)
# For now, this might succeed (no validation)
# After adding validation, we'd expect a ValidationError
except Exception:
# Expected once validation is implemented
pass
class TestOutputValidation:
"""Test output schema validation (future feature)."""
def test_output_matches_schema(self, llama_stack_client, mcp_server_with_output_schemas):
"""Test that tool output is validated against output_schema (if implemented)."""
if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
pytest.skip("Library client required for local MCP server")
test_toolgroup_id = "mcp::with_output"
uri = mcp_server_with_output_schemas["server_url"]
try:
llama_stack_client.toolgroups.unregister(toolgroup_id=test_toolgroup_id)
except Exception:
pass
llama_stack_client.toolgroups.register(
toolgroup_id=test_toolgroup_id,
provider_id="model-context-protocol",
mcp_endpoint=dict(uri=uri),
)
provider_data = {"mcp_headers": {uri: {"Authorization": f"Bearer {AUTH_TOKEN}"}}}
auth_headers = {
"X-LlamaStack-Provider-Data": json.dumps(provider_data),
}
# Invoke tool
result = llama_stack_client.tool_runtime.invoke_tool(
tool_name="get_weather",
kwargs={"location": "San Francisco"},
extra_headers=auth_headers,
)
# Tool should return valid output
assert result.error_message is None
assert result.content is not None
# Once output validation is implemented, the system would check
# that result.content matches the tool's output_schema

76
tests/unit/models/test_prompt_adapter.py
View file

@ -18,7 +18,6 @@ from llama_stack.apis.inference import (
from llama_stack.models.llama.datatypes import (
BuiltinTool,
ToolDefinition,
ToolParamDefinition,
ToolPromptFormat,
)
from llama_stack.providers.utils.inference.prompt_adapter import (
@ -75,12 +74,15 @@ async def test_system_custom_only():
ToolDefinition(
tool_name="custom1",
description="custom1 tool",
parameters={
"param1": ToolParamDefinition(
param_type="str",
description="param1 description",
required=True,
),
input_schema={
"type": "object",
"properties": {
"param1": {
"type": "str",
"description": "param1 description",
},
},
"required": ["param1"],
},
)
],
@ -107,12 +109,15 @@ async def test_system_custom_and_builtin():
ToolDefinition(
tool_name="custom1",
description="custom1 tool",
parameters={
"param1": ToolParamDefinition(
param_type="str",
description="param1 description",
required=True,
),
input_schema={
"type": "object",
"properties": {
"param1": {
"type": "str",
"description": "param1 description",
},
},
"required": ["param1"],
},
),
],
@ -148,12 +153,15 @@ async def test_completion_message_encoding():
ToolDefinition(
tool_name="custom1",
description="custom1 tool",
parameters={
"param1": ToolParamDefinition(
param_type="str",
description="param1 description",
required=True,
),
input_schema={
"type": "object",
"properties": {
"param1": {
"type": "str",
"description": "param1 description",
},
},
"required": ["param1"],
},
),
],
@ -227,12 +235,15 @@ async def test_replace_system_message_behavior_custom_tools():
ToolDefinition(
tool_name="custom1",
description="custom1 tool",
parameters={
"param1": ToolParamDefinition(
param_type="str",
description="param1 description",
required=True,
),
input_schema={
"type": "object",
"properties": {
"param1": {
"type": "str",
"description": "param1 description",
},
},
"required": ["param1"],
},
),
],
@ -264,12 +275,15 @@ async def test_replace_system_message_behavior_custom_tools_with_template():
ToolDefinition(
tool_name="custom1",
description="custom1 tool",
parameters={
"param1": ToolParamDefinition(
param_type="str",
description="param1 description",
required=True,
),
input_schema={
"type": "object",
"properties": {
"param1": {
"type": "str",
"description": "param1 description",
},
},
"required": ["param1"],
},
),
],

381
tests/unit/providers/utils/test_openai_compat_conversion.py Normal file
View file

@ -0,0 +1,381 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
"""
Unit tests for OpenAI compatibility tool conversion.
Tests convert_tooldef_to_openai_tool with new JSON Schema approach.
"""
from llama_stack.models.llama.datatypes import BuiltinTool, ToolDefinition
from llama_stack.providers.utils.inference.openai_compat import convert_tooldef_to_openai_tool
class TestSimpleSchemaConversion:
"""Test basic schema conversions to OpenAI format."""
def test_simple_tool_conversion(self):
"""Test conversion of simple tool with basic input schema."""
tool = ToolDefinition(
tool_name="get_weather",
description="Get weather information",
input_schema={
"type": "object",
"properties": {"location": {"type": "string", "description": "City name"}},
"required": ["location"],
},
)
result = convert_tooldef_to_openai_tool(tool)
# Check OpenAI structure
assert result["type"] == "function"
assert "function" in result
function = result["function"]
assert function["name"] == "get_weather"
assert function["description"] == "Get weather information"
# Check parameters are passed through
assert "parameters" in function
assert function["parameters"] == tool.input_schema
assert function["parameters"]["type"] == "object"
assert "location" in function["parameters"]["properties"]
def test_tool_without_description(self):
"""Test tool conversion without description."""
tool = ToolDefinition(tool_name="test_tool", input_schema={"type": "object", "properties": {}})
result = convert_tooldef_to_openai_tool(tool)
assert result["function"]["name"] == "test_tool"
assert "description" not in result["function"]
assert "parameters" in result["function"]
def test_builtin_tool_conversion(self):
"""Test conversion of BuiltinTool enum."""
tool = ToolDefinition(
tool_name=BuiltinTool.code_interpreter,
description="Run Python code",
input_schema={"type": "object", "properties": {"code": {"type": "string"}}},
)
result = convert_tooldef_to_openai_tool(tool)
# BuiltinTool should be converted to its value
assert result["function"]["name"] == "code_interpreter"
class TestComplexSchemaConversion:
"""Test conversion of complex JSON Schema features."""
def test_schema_with_refs_and_defs(self):
"""Test that $ref and $defs are passed through to OpenAI."""
tool = ToolDefinition(
tool_name="book_flight",
description="Book a flight",
input_schema={
"type": "object",
"properties": {
"flight": {"$ref": "#/$defs/FlightInfo"},
"passengers": {"type": "array", "items": {"$ref": "#/$defs/Passenger"}},
"payment": {"$ref": "#/$defs/Payment"},
},
"required": ["flight", "passengers", "payment"],
"$defs": {
"FlightInfo": {
"type": "object",
"properties": {
"from": {"type": "string", "description": "Departure airport"},
"to": {"type": "string", "description": "Arrival airport"},
"date": {"type": "string", "format": "date"},
},
"required": ["from", "to", "date"],
},
"Passenger": {
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer", "minimum": 0}},
"required": ["name", "age"],
},
"Payment": {
"type": "object",
"properties": {
"method": {"type": "string", "enum": ["credit_card", "debit_card"]},
"amount": {"type": "number", "minimum": 0},
},
},
},
},
)
result = convert_tooldef_to_openai_tool(tool)
params = result["function"]["parameters"]
# Verify $defs are preserved
assert "$defs" in params
assert "FlightInfo" in params["$defs"]
assert "Passenger" in params["$defs"]
assert "Payment" in params["$defs"]
# Verify $ref are preserved
assert params["properties"]["flight"]["$ref"] == "#/$defs/FlightInfo"
assert params["properties"]["passengers"]["items"]["$ref"] == "#/$defs/Passenger"
assert params["properties"]["payment"]["$ref"] == "#/$defs/Payment"
# Verify nested schema details are preserved
assert params["$defs"]["FlightInfo"]["properties"]["date"]["format"] == "date"
assert params["$defs"]["Passenger"]["properties"]["age"]["minimum"] == 0
assert params["$defs"]["Payment"]["properties"]["method"]["enum"] == ["credit_card", "debit_card"]
def test_anyof_schema_conversion(self):
"""Test conversion of anyOf schemas."""
tool = ToolDefinition(
tool_name="flexible_input",
input_schema={
"type": "object",
"properties": {
"contact": {
"anyOf": [
{"type": "string", "format": "email"},
{"type": "string", "pattern": "^\\+?[0-9]{10,15}$"},
],
"description": "Email or phone number",
}
},
},
)
result = convert_tooldef_to_openai_tool(tool)
contact_schema = result["function"]["parameters"]["properties"]["contact"]
assert "anyOf" in contact_schema
assert len(contact_schema["anyOf"]) == 2
assert contact_schema["anyOf"][0]["format"] == "email"
assert "pattern" in contact_schema["anyOf"][1]
def test_nested_objects_conversion(self):
"""Test conversion of deeply nested objects."""
tool = ToolDefinition(
tool_name="nested_data",
input_schema={
"type": "object",
"properties": {
"user": {
"type": "object",
"properties": {
"profile": {
"type": "object",
"properties": {
"name": {"type": "string"},
"settings": {
"type": "object",
"properties": {"theme": {"type": "string", "enum": ["light", "dark"]}},
},
},
}
},
}
},
},
)
result = convert_tooldef_to_openai_tool(tool)
# Navigate deep structure
user_schema = result["function"]["parameters"]["properties"]["user"]
profile_schema = user_schema["properties"]["profile"]
settings_schema = profile_schema["properties"]["settings"]
theme_schema = settings_schema["properties"]["theme"]
assert theme_schema["enum"] == ["light", "dark"]
def test_array_schemas_with_constraints(self):
"""Test conversion of array schemas with constraints."""
tool = ToolDefinition(
tool_name="list_processor",
input_schema={
"type": "object",
"properties": {
"items": {
"type": "array",
"items": {
"type": "object",
"properties": {"id": {"type": "integer"}, "name": {"type": "string"}},
"required": ["id"],
},
"minItems": 1,
"maxItems": 100,
"uniqueItems": True,
}
},
},
)
result = convert_tooldef_to_openai_tool(tool)
items_schema = result["function"]["parameters"]["properties"]["items"]
assert items_schema["type"] == "array"
assert items_schema["minItems"] == 1
assert items_schema["maxItems"] == 100
assert items_schema["uniqueItems"] is True
assert items_schema["items"]["type"] == "object"
class TestOutputSchemaHandling:
"""Test that output_schema is correctly handled (or dropped) for OpenAI."""
def test_output_schema_is_dropped(self):
"""Test that output_schema is NOT included in OpenAI format (API limitation)."""
tool = ToolDefinition(
tool_name="calculator",
description="Perform calculation",
input_schema={"type": "object", "properties": {"x": {"type": "number"}, "y": {"type": "number"}}},
output_schema={"type": "object", "properties": {"result": {"type": "number"}}, "required": ["result"]},
)
result = convert_tooldef_to_openai_tool(tool)
# OpenAI doesn't support output schema
assert "outputSchema" not in result["function"]
assert "responseSchema" not in result["function"]
assert "output_schema" not in result["function"]
# But input schema should be present
assert "parameters" in result["function"]
assert result["function"]["parameters"] == tool.input_schema
def test_only_output_schema_no_input(self):
"""Test tool with only output_schema (unusual but valid)."""
tool = ToolDefinition(
tool_name="no_input_tool",
description="Tool with no inputs",
output_schema={"type": "object", "properties": {"timestamp": {"type": "string"}}},
)
result = convert_tooldef_to_openai_tool(tool)
# No parameters should be set if input_schema is None
# (or we might set an empty object schema - implementation detail)
assert "outputSchema" not in result["function"]
class TestEdgeCases:
"""Test edge cases and error conditions."""
def test_tool_with_no_schemas(self):
"""Test tool with neither input nor output schema."""
tool = ToolDefinition(tool_name="schemaless_tool", description="Tool without schemas")
result = convert_tooldef_to_openai_tool(tool)
assert result["function"]["name"] == "schemaless_tool"
assert result["function"]["description"] == "Tool without schemas"
# Implementation detail: might have no parameters or empty object
def test_empty_input_schema(self):
"""Test tool with empty object schema."""
tool = ToolDefinition(tool_name="no_params", input_schema={"type": "object", "properties": {}})
result = convert_tooldef_to_openai_tool(tool)
assert result["function"]["parameters"]["type"] == "object"
assert result["function"]["parameters"]["properties"] == {}
def test_schema_with_additional_properties(self):
"""Test that additionalProperties is preserved."""
tool = ToolDefinition(
tool_name="flexible_tool",
input_schema={
"type": "object",
"properties": {"known_field": {"type": "string"}},
"additionalProperties": True,
},
)
result = convert_tooldef_to_openai_tool(tool)
assert result["function"]["parameters"]["additionalProperties"] is True
def test_schema_with_pattern_properties(self):
"""Test that patternProperties is preserved."""
tool = ToolDefinition(
tool_name="pattern_tool",
input_schema={"type": "object", "patternProperties": {"^[a-z]+$": {"type": "string"}}},
)
result = convert_tooldef_to_openai_tool(tool)
assert "patternProperties" in result["function"]["parameters"]
def test_schema_identity(self):
"""Test that converted schema is identical to input (no lossy conversion)."""
original_schema = {
"type": "object",
"properties": {"complex": {"$ref": "#/$defs/Complex"}},
"$defs": {
"Complex": {
"type": "object",
"properties": {"nested": {"anyOf": [{"type": "string"}, {"type": "number"}]}},
}
},
"required": ["complex"],
"additionalProperties": False,
}
tool = ToolDefinition(tool_name="test", input_schema=original_schema)
result = convert_tooldef_to_openai_tool(tool)
# Converted parameters should be EXACTLY the same as input
assert result["function"]["parameters"] == original_schema
class TestConversionConsistency:
"""Test consistency across multiple conversions."""
def test_multiple_tools_with_shared_defs(self):
"""Test converting multiple tools that could share definitions."""
tool1 = ToolDefinition(
tool_name="tool1",
input_schema={
"type": "object",
"properties": {"data": {"$ref": "#/$defs/Data"}},
"$defs": {"Data": {"type": "object", "properties": {"x": {"type": "number"}}}},
},
)
tool2 = ToolDefinition(
tool_name="tool2",
input_schema={
"type": "object",
"properties": {"info": {"$ref": "#/$defs/Data"}},
"$defs": {"Data": {"type": "object", "properties": {"y": {"type": "string"}}}},
},
)
result1 = convert_tooldef_to_openai_tool(tool1)
result2 = convert_tooldef_to_openai_tool(tool2)
# Each tool maintains its own $defs independently
assert result1["function"]["parameters"]["$defs"]["Data"]["properties"]["x"]["type"] == "number"
assert result2["function"]["parameters"]["$defs"]["Data"]["properties"]["y"]["type"] == "string"
def test_conversion_is_pure(self):
"""Test that conversion doesn't modify the original tool."""
original_schema = {
"type": "object",
"properties": {"x": {"type": "string"}},
"$defs": {"T": {"type": "number"}},
}
tool = ToolDefinition(tool_name="test", input_schema=original_schema.copy())
# Convert
convert_tooldef_to_openai_tool(tool)
# Original tool should be unchanged
assert tool.input_schema == original_schema
assert "$defs" in tool.input_schema

297
tests/unit/tools/test_tools_json_schema.py Normal file
View file

@ -0,0 +1,297 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
"""
Unit tests for JSON Schema-based tool definitions.
Tests the new input_schema and output_schema fields.
"""
from pydantic import ValidationError
from llama_stack.apis.tools import ToolDef
from llama_stack.models.llama.datatypes import BuiltinTool, ToolDefinition
class TestToolDefValidation:
"""Test ToolDef validation with JSON Schema."""
def test_simple_input_schema(self):
"""Test ToolDef with simple input schema."""
tool = ToolDef(
name="get_weather",
description="Get weather information",
input_schema={
"type": "object",
"properties": {"location": {"type": "string", "description": "City name"}},
"required": ["location"],
},
)
assert tool.name == "get_weather"
assert tool.input_schema["type"] == "object"
assert "location" in tool.input_schema["properties"]
assert tool.output_schema is None
def test_input_and_output_schema(self):
"""Test ToolDef with both input and output schemas."""
tool = ToolDef(
name="calculate",
description="Perform calculation",
input_schema={
"type": "object",
"properties": {"x": {"type": "number"}, "y": {"type": "number"}},
"required": ["x", "y"],
},
output_schema={"type": "object", "properties": {"result": {"type": "number"}}, "required": ["result"]},
)
assert tool.input_schema is not None
assert tool.output_schema is not None
assert "result" in tool.output_schema["properties"]
def test_schema_with_refs_and_defs(self):
"""Test that $ref and $defs are preserved in schemas."""
tool = ToolDef(
name="book_flight",
description="Book a flight",
input_schema={
"type": "object",
"properties": {
"flight": {"$ref": "#/$defs/FlightInfo"},
"passengers": {"type": "array", "items": {"$ref": "#/$defs/Passenger"}},
},
"$defs": {
"FlightInfo": {
"type": "object",
"properties": {"from": {"type": "string"}, "to": {"type": "string"}},
},
"Passenger": {
"type": "object",
"properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
},
},
},
)
# Verify $defs are preserved
assert "$defs" in tool.input_schema
assert "FlightInfo" in tool.input_schema["$defs"]
assert "Passenger" in tool.input_schema["$defs"]
# Verify $ref are preserved
assert tool.input_schema["properties"]["flight"]["$ref"] == "#/$defs/FlightInfo"
assert tool.input_schema["properties"]["passengers"]["items"]["$ref"] == "#/$defs/Passenger"
def test_output_schema_with_refs(self):
"""Test that output_schema also supports $ref and $defs."""
tool = ToolDef(
name="search",
description="Search for items",
input_schema={"type": "object", "properties": {"query": {"type": "string"}}},
output_schema={
"type": "object",
"properties": {"results": {"type": "array", "items": {"$ref": "#/$defs/SearchResult"}}},
"$defs": {
"SearchResult": {
"type": "object",
"properties": {"title": {"type": "string"}, "score": {"type": "number"}},
}
},
},
)
assert "$defs" in tool.output_schema
assert "SearchResult" in tool.output_schema["$defs"]
def test_complex_json_schema_features(self):
"""Test various JSON Schema features are preserved."""
tool = ToolDef(
name="complex_tool",
description="Tool with complex schema",
input_schema={
"type": "object",
"properties": {
# anyOf
"contact": {
"anyOf": [
{"type": "string", "format": "email"},
{"type": "string", "pattern": "^\\+?[0-9]{10,15}$"},
]
},
# enum
"status": {"type": "string", "enum": ["pending", "approved", "rejected"]},
# nested objects
"address": {
"type": "object",
"properties": {
"street": {"type": "string"},
"city": {"type": "string"},
"zipcode": {"type": "string", "pattern": "^[0-9]{5}$"},
},
"required": ["street", "city"],
},
# array with constraints
"tags": {
"type": "array",
"items": {"type": "string"},
"minItems": 1,
"maxItems": 10,
"uniqueItems": True,
},
},
},
)
# Verify anyOf
assert "anyOf" in tool.input_schema["properties"]["contact"]
# Verify enum
assert tool.input_schema["properties"]["status"]["enum"] == ["pending", "approved", "rejected"]
# Verify nested object
assert tool.input_schema["properties"]["address"]["type"] == "object"
assert "zipcode" in tool.input_schema["properties"]["address"]["properties"]
# Verify array constraints
tags_schema = tool.input_schema["properties"]["tags"]
assert tags_schema["minItems"] == 1
assert tags_schema["maxItems"] == 10
assert tags_schema["uniqueItems"] is True
def test_invalid_json_schema_raises_error(self):
"""Test that invalid JSON Schema raises validation error."""
# TODO: This test will pass once we add schema validation
# For now, Pydantic accepts any dict, so this is a placeholder
# This should eventually raise an error due to invalid schema
try:
ToolDef(
name="bad_tool",
input_schema={
"type": "invalid_type", # Not a valid JSON Schema type
"properties": "not_an_object", # Should be an object
},
)
# For now this passes, but shouldn't after we add validation
except ValidationError:
pass # Expected once validation is added
class TestToolDefinitionValidation:
"""Test ToolDefinition (internal) validation with JSON Schema."""
def test_simple_tool_definition(self):
"""Test ToolDefinition with simple schema."""
tool = ToolDefinition(
tool_name="get_time",
description="Get current time",
input_schema={"type": "object", "properties": {"timezone": {"type": "string"}}},
)
assert tool.tool_name == "get_time"
assert tool.input_schema is not None
def test_builtin_tool_with_schema(self):
"""Test ToolDefinition with BuiltinTool enum."""
tool = ToolDefinition(
tool_name=BuiltinTool.code_interpreter,
description="Run Python code",
input_schema={"type": "object", "properties": {"code": {"type": "string"}}, "required": ["code"]},
output_schema={"type": "object", "properties": {"output": {"type": "string"}, "error": {"type": "string"}}},
)
assert isinstance(tool.tool_name, BuiltinTool)
assert tool.input_schema is not None
assert tool.output_schema is not None
def test_tool_definition_with_refs(self):
"""Test ToolDefinition preserves $ref/$defs."""
tool = ToolDefinition(
tool_name="process_data",
input_schema={
"type": "object",
"properties": {"data": {"$ref": "#/$defs/DataObject"}},
"$defs": {
"DataObject": {
"type": "object",
"properties": {
"id": {"type": "integer"},
"values": {"type": "array", "items": {"type": "number"}},
},
}
},
},
)
assert "$defs" in tool.input_schema
assert tool.input_schema["properties"]["data"]["$ref"] == "#/$defs/DataObject"
class TestSchemaEquivalence:
"""Test that schemas remain unchanged through serialization."""
def test_schema_roundtrip(self):
"""Test that schemas survive model_dump/model_validate roundtrip."""
original = ToolDef(
name="test",
input_schema={
"type": "object",
"properties": {"x": {"$ref": "#/$defs/X"}},
"$defs": {"X": {"type": "string"}},
},
)
# Serialize and deserialize
dumped = original.model_dump()
restored = ToolDef(**dumped)
# Schemas should be identical
assert restored.input_schema == original.input_schema
assert "$defs" in restored.input_schema
assert restored.input_schema["properties"]["x"]["$ref"] == "#/$defs/X"
def test_json_serialization(self):
"""Test JSON serialization preserves schema."""
import json
tool = ToolDef(
name="test",
input_schema={
"type": "object",
"properties": {"a": {"type": "string"}},
"$defs": {"T": {"type": "number"}},
},
output_schema={"type": "object", "properties": {"b": {"$ref": "#/$defs/T"}}},
)
# Serialize to JSON and back
json_str = tool.model_dump_json()
parsed = json.loads(json_str)
restored = ToolDef(**parsed)
assert restored.input_schema == tool.input_schema
assert restored.output_schema == tool.output_schema
assert "$defs" in restored.input_schema
class TestBackwardsCompatibility:
"""Test handling of legacy code patterns."""
def test_none_schemas(self):
"""Test tools with no schemas (legacy case)."""
tool = ToolDef(name="legacy_tool", description="Tool without schemas", input_schema=None, output_schema=None)
assert tool.input_schema is None
assert tool.output_schema is None
def test_metadata_preserved(self):
"""Test that metadata field still works."""
tool = ToolDef(
name="test", input_schema={"type": "object"}, metadata={"endpoint": "http://example.com", "version": "1.0"}
)
assert tool.metadata["endpoint"] == "http://example.com"
assert tool.metadata["version"] == "1.0"