phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 12:07:34 +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

12
llama_stack/models/llama/datatypes.py
View file

@ -88,19 +88,11 @@ class StopReason(Enum):
out_of_tokens = "out_of_tokens"
class ToolParamDefinition(BaseModel):
param_type: str
description: str | None = None
required: bool | None = True
items: Any | None = None
title: str | None = None
default: Any | None = None
class ToolDefinition(BaseModel):
tool_name: BuiltinTool | str
description: str | None = None
parameters: dict[str, ToolParamDefinition] | None = None
input_schema: dict[str, Any] | None = None
output_schema: dict[str, Any] | None = None
@field_validator("tool_name", mode="before")
@classmethod

112
llama_stack/models/llama/llama3/prompt_templates/system_prompts.py
View file

@ -18,7 +18,6 @@ from typing import Any
from llama_stack.apis.inference import (
BuiltinTool,
ToolDefinition,
ToolParamDefinition,
)
from .base import PromptTemplate, PromptTemplateGeneratorBase
@ -101,11 +100,8 @@ class JsonCustomToolGenerator(PromptTemplateGeneratorBase):
{# manually setting up JSON because jinja sorts keys in unexpected ways -#}
{%- set tname = t.tool_name -%}
{%- set tdesc = t.description -%}
{%- set tparams = t.parameters -%}
{%- set required_params = [] -%}
{%- for name, param in tparams.items() if param.required == true -%}
{%- set _ = required_params.append(name) -%}
{%- endfor -%}
{%- set tprops = t.input_schema.get('properties', {}) -%}
{%- set required_params = t.input_schema.get('required', []) -%}
{
"type": "function",
"function": {
@ -114,11 +110,11 @@ class JsonCustomToolGenerator(PromptTemplateGeneratorBase):
"parameters": {
"type": "object",
"properties": [
{%- for name, param in tparams.items() %}
{%- for name, param in tprops.items() %}
{
"{{name}}": {
"type": "object",
"description": "{{param.description}}"
"description": "{{param.get('description', '')}}"
}
}{% if not loop.last %},{% endif %}
{%- endfor %}
@ -143,17 +139,19 @@ class JsonCustomToolGenerator(PromptTemplateGeneratorBase):
ToolDefinition(
tool_name="trending_songs",
description="Returns the trending songs on a Music site",
parameters={
"n": ToolParamDefinition(
param_type="int",
description="The number of songs to return",
required=True,
),
"genre": ToolParamDefinition(
param_type="str",
description="The genre of the songs to return",
required=False,
),
input_schema={
"type": "object",
"properties": {
"n": {
"type": "int",
"description": "The number of songs to return",
},
"genre": {
"type": "str",
"description": "The genre of the songs to return",
},
},
"required": ["n"],
},
),
]
@ -170,11 +168,14 @@ class FunctionTagCustomToolGenerator(PromptTemplateGeneratorBase):
{#- manually setting up JSON because jinja sorts keys in unexpected ways -#}
{%- set tname = t.tool_name -%}
{%- set tdesc = t.description -%}
{%- set modified_params = t.parameters.copy() -%}
{%- for key, value in modified_params.items() -%}
{%- if 'default' in value -%}
{%- set _ = value.pop('default', None) -%}
{%- set tprops = t.input_schema.get('properties', {}) -%}
{%- set modified_params = {} -%}
{%- for key, value in tprops.items() -%}
{%- set param_copy = value.copy() -%}
{%- if 'default' in param_copy -%}
{%- set _ = param_copy.pop('default', None) -%}
{%- endif -%}
{%- set _ = modified_params.update({key: param_copy}) -%}
{%- endfor -%}
{%- set tparams = modified_params | tojson -%}
Use the function '{{ tname }}' to '{{ tdesc }}':
@ -205,17 +206,19 @@ class FunctionTagCustomToolGenerator(PromptTemplateGeneratorBase):
ToolDefinition(
tool_name="trending_songs",
description="Returns the trending songs on a Music site",
parameters={
"n": ToolParamDefinition(
param_type="int",
description="The number of songs to return",
required=True,
),
"genre": ToolParamDefinition(
param_type="str",
description="The genre of the songs to return",
required=False,
),
input_schema={
"type": "object",
"properties": {
"n": {
"type": "int",
"description": "The number of songs to return",
},
"genre": {
"type": "str",
"description": "The genre of the songs to return",
},
},
"required": ["n"],
},
),
]
@ -255,11 +258,8 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase): # noqa: N801
{# manually setting up JSON because jinja sorts keys in unexpected ways -#}
{%- set tname = t.tool_name -%}
{%- set tdesc = t.description -%}
{%- set tparams = t.parameters -%}
{%- set required_params = [] -%}
{%- for name, param in tparams.items() if param.required == true -%}
{%- set _ = required_params.append(name) -%}
{%- endfor -%}
{%- set tprops = t.input_schema.get('properties', {}) -%}
{%- set required_params = t.input_schema.get('required', []) -%}
{
"name": "{{tname}}",
"description": "{{tdesc}}",
@ -267,11 +267,11 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase): # noqa: N801
"type": "dict",
"required": {{ required_params | tojson }},
"properties": {
{%- for name, param in tparams.items() %}
{%- for name, param in tprops.items() %}
"{{name}}": {
"type": "{{param.param_type}}",
"description": "{{param.description}}"{% if param.default %},
"default": "{{param.default}}"{% endif %}
"type": "{{param.get('type', 'string')}}",
"description": "{{param.get('description', '')}}"{% if param.get('default') %},
"default": "{{param.get('default')}}"{% endif %}
}{% if not loop.last %},{% endif %}
{%- endfor %}
}
@ -299,18 +299,20 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase): # noqa: N801
ToolDefinition(
tool_name="get_weather",
description="Get weather info for places",
parameters={
"city": ToolParamDefinition(
param_type="string",
description="The name of the city to get the weather for",
required=True,
),
"metric": ToolParamDefinition(
param_type="string",
description="The metric for weather. Options are: celsius, fahrenheit",
required=False,
default="celsius",
),
input_schema={
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The name of the city to get the weather for",
},
"metric": {
"type": "string",
"description": "The metric for weather. Options are: celsius, fahrenheit",
"default": "celsius",
},
},
"required": ["city"],
},
),
]

43
llama_stack/models/llama/llama4/prompt_templates/system_prompts.py
View file

@ -13,7 +13,7 @@
import textwrap
from llama_stack.apis.inference import ToolDefinition, ToolParamDefinition
from llama_stack.apis.inference import ToolDefinition
from llama_stack.models.llama.llama3.prompt_templates.base import (
PromptTemplate,
PromptTemplateGeneratorBase,
@ -81,11 +81,8 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase): # noqa: N801
{# manually setting up JSON because jinja sorts keys in unexpected ways -#}
{%- set tname = t.tool_name -%}
{%- set tdesc = t.description -%}
{%- set tparams = t.parameters -%}
{%- set required_params = [] -%}
{%- for name, param in tparams.items() if param.required == true -%}
{%- set _ = required_params.append(name) -%}
{%- endfor -%}
{%- set tprops = t.input_schema.get('properties', {}) -%}
{%- set required_params = t.input_schema.get('required', []) -%}
{
"name": "{{tname}}",
"description": "{{tdesc}}",
@ -93,11 +90,11 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase): # noqa: N801
"type": "dict",
"required": {{ required_params | tojson }},
"properties": {
{%- for name, param in tparams.items() %}
{%- for name, param in tprops.items() %}
"{{name}}": {
"type": "{{param.param_type}}",
"description": "{{param.description}}"{% if param.default %},
"default": "{{param.default}}"{% endif %}
"type": "{{param.get('type', 'string')}}",
"description": "{{param.get('description', '')}}"{% if param.get('default') %},
"default": "{{param.get('default')}}"{% endif %}
}{% if not loop.last %},{% endif %}
{%- endfor %}
}
@ -119,18 +116,20 @@ class PythonListCustomToolGenerator(PromptTemplateGeneratorBase): # noqa: N801
ToolDefinition(
tool_name="get_weather",
description="Get weather info for places",
parameters={
"city": ToolParamDefinition(
param_type="string",
description="The name of the city to get the weather for",
required=True,
),
"metric": ToolParamDefinition(
param_type="string",
description="The metric for weather. Options are: celsius, fahrenheit",
required=False,
default="celsius",
),
input_schema={
"type": "object",
"properties": {
"city": {
"type": "string",
"description": "The name of the city to get the weather for",
},
"metric": {
"type": "string",
"description": "The metric for weather. Options are: celsius, fahrenheit",
"default": "celsius",
},
},
"required": ["city"],
},
),
]