use super::super::tool::{
field::{Choice, FieldSchemaAttr},
json::JsonType,
};
use crate::resolve;
use crate::schema_emit;
use proc_macro::TokenStream;
use quote::quote;
use serde::Serialize;
use std::collections::BTreeMap;
use strum::{Display, EnumString};
use syn::{
Attribute, Data, DataStruct, DeriveInput, Error, Field, Ident, LitStr, Result, Type,
parse_macro_input, spanned::Spanned,
};
#[derive(EnumString, Display)]
enum OutputAttrIdent {
#[strum(serialize = "output")]
Output,
}
#[derive(Debug, Serialize)]
pub(crate) struct OutputSchemaProperty {
#[serde(rename = "type")]
_type: String,
description: Option<String>,
#[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
_enum: Option<Vec<serde_json::Value>>,
}
#[derive(Debug, Serialize, Default)]
pub(crate) struct OutputSchema {
#[serde(rename = "type")]
_type: String,
#[serde(default)]
properties: BTreeMap<String, OutputSchemaProperty>,
#[serde(default)]
required: Vec<String>,
}
#[derive(Debug, Serialize, Default)]
pub(crate) struct StructuredOutputFormat {
name: String,
description: Option<String>,
#[serde(default)]
schema: OutputSchema,
strict: Option<bool>,
}
#[derive(Debug, Default)]
pub(crate) struct OutputParser {
output_data: StructuredOutputFormat,
ident: Option<Ident>,
}
impl OutputParser {
pub fn parse(&mut self, input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let struct_ident = input.ident.clone();
let struct_span = struct_ident.span();
self.ident = Some(input.ident);
self.output_data.name = struct_ident.to_string();
self.output_data.schema._type = JsonType::Object.to_string();
if let Err(err) = self.parse_struct_attributes(&input.attrs) {
return err.to_compile_error().into();
}
if let Err(err) = self.parse_data(input.data) {
return err.to_compile_error().into();
}
let core = match resolve::resolve_core_path() {
Ok(core) => core,
Err(err) => return err.to_compile_error().into(),
};
let core = &core;
let serialized_data = match serde_json::to_string(&self.output_data) {
Ok(data) => data,
Err(err) => {
return Error::new(
struct_span,
format!("failed to serialize agent output schema: {err}"),
)
.to_compile_error()
.into();
}
};
let schema_tokens = match schema_emit::schema_str_to_tokens(&serialized_data, struct_span) {
Ok(tokens) => tokens,
Err(err) => return err.to_compile_error().into(),
};
let schema_literal = LitStr::new(&serialized_data, struct_span);
let expanded = quote! {
impl #core::agent::AgentOutputT for #struct_ident {
fn output_schema() -> &'static str {
#schema_literal
}
fn structured_output_format() -> ::serde_json::Value {
static SCHEMA: ::std::sync::LazyLock<::serde_json::Value> =
::std::sync::LazyLock::new(|| #schema_tokens);
(*SCHEMA).clone()
}
}
};
TokenStream::from(expanded)
}
fn parse_struct_attributes(&mut self, attrs: &[Attribute]) -> Result<()> {
for attr in attrs {
if attr.path().is_ident("doc") {
if let syn::Meta::NameValue(meta) = &attr.meta
&& let syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Str(lit_str),
..
}) = &meta.value
{
let doc_value = lit_str.value().trim().to_string();
if !doc_value.is_empty() {
let description =
self.output_data.description.get_or_insert_with(String::new);
if !description.is_empty() {
description.push(' ');
}
description.push_str(&doc_value);
}
}
} else if attr.path().is_ident("strict") {
if let Ok(strict_value) = attr.parse_args::<syn::LitBool>() {
self.output_data.strict = Some(strict_value.value);
} else {
return Err(Error::new(
attr.span(),
"`#[strict]` on AgentOutput structs must be a boolean literal, e.g. `#[strict(true)]`",
));
}
}
}
Ok(())
}
fn parse_data(&mut self, input: Data) -> Result<()> {
match &input {
Data::Struct(struct_data) => self.parse_struct(struct_data)?,
_ => {
return Err(Error::new(
proc_macro2::Span::call_site(),
"Union or Enums not yet supported!",
));
}
};
Ok(())
}
fn parse_struct(&mut self, input: &DataStruct) -> Result<()> {
match &input.fields {
syn::Fields::Named(fields) => {
let mut has_output_attribute = false;
for field in fields.named.iter() {
let field_name = field.ident.as_ref().ok_or_else(|| {
Error::new(field.span(), "named fields must have an identifier")
})?;
let field_name = field_name.to_string();
let has_field_output_attr = field.attrs.iter().any(|attr| {
attr.path()
.is_ident(OutputAttrIdent::Output.to_string().as_str())
});
if has_field_output_attr {
has_output_attribute = true;
}
let output_property = self.parse_field(field_name.clone(), field)?;
self.output_data
.schema
.properties
.insert(field_name, output_property);
}
if !has_output_attribute {
return Err(Error::new(
proc_macro2::Span::call_site(),
"AgentOutput structs must have at least one field with an #[output(description = \"...\")] attribute",
));
}
}
_ => {
return Err(Error::new(
proc_macro2::Span::call_site(),
"Tuple or Unit structs not yet supported!",
));
}
}
Ok(())
}
fn parse_field(&mut self, name: String, field: &Field) -> Result<OutputSchemaProperty> {
let (is_optional, inner_type) = self.extract_option_type(&field.ty);
if !is_optional {
self.output_data.schema.required.push(name.clone());
}
let json_type = self.get_json_type(inner_type.unwrap_or(&field.ty))?;
let mut field_schema: Option<FieldSchemaAttr> = None;
for attr in &field.attrs {
if attr
.path()
.is_ident(OutputAttrIdent::Output.to_string().as_str())
{
field_schema = Some(self.parse_field_attributes(attr, &json_type)?);
}
}
if let Some(schema) = field_schema {
Ok(OutputSchemaProperty {
_type: json_type.to_string(),
description: schema.description.map(|lit| lit.value()),
_enum: schema.choice.map(choices_to_json_values).transpose()?,
})
} else {
Ok(OutputSchemaProperty {
_type: json_type.to_string(),
description: None,
_enum: None,
})
}
}
fn extract_option_type<'a>(&self, ty: &'a Type) -> (bool, Option<&'a Type>) {
if let Type::Path(type_path) = ty
&& let Some(segment) = type_path.path.segments.last()
&& segment.ident == "Option"
&& let syn::PathArguments::AngleBracketed(args) = &segment.arguments
&& let Some(syn::GenericArgument::Type(inner_type)) = args.args.first()
{
return (true, Some(inner_type));
}
(false, None)
}
fn get_json_type(&self, field_type: &Type) -> Result<JsonType> {
match field_type {
Type::Path(path) => {
let Some(segment) = path.path.segments.last() else {
return Err(Error::new(
proc_macro2::Span::call_site(),
"Invalid type path in AgentOutput field",
));
};
if segment.ident == "Vec" {
return Ok(JsonType::Array);
}
match segment.ident.to_string().as_str() {
"String" | "str" => Ok(JsonType::String),
"i8" | "i32" | "u32" | "u8" | "i64" | "u64" | "i16" | "u16" | "isize"
| "usize" => Ok(JsonType::Integer),
"f64" | "f32" => Ok(JsonType::Number),
"bool" => Ok(JsonType::Boolean),
other => Err(Error::new(
proc_macro2::Span::call_site(),
format!("Unsupported data type: {other}"),
)),
}
}
Type::Reference(reference) => self.get_json_type(&reference.elem),
Type::Group(group) => self.get_json_type(&group.elem),
Type::Paren(paren) => self.get_json_type(&paren.elem),
other => Err(Error::new(
proc_macro2::Span::call_site(),
format!("Unsupported AgentOutput field type: {other:?}"),
)),
}
}
fn parse_field_attributes(
&self,
attribute: &Attribute,
field_type: &JsonType,
) -> Result<FieldSchemaAttr> {
let attributes = attribute.parse_args::<FieldSchemaAttr>()?;
if let Some(ref enum_vals) = attributes.choice {
let invalid_choice = enum_vals.iter().find(|c| {
!matches!(
(c, field_type),
(Choice::String(_), JsonType::String)
| (Choice::Number(_), JsonType::Number)
| (Choice::Number(_), JsonType::Integer)
)
});
if invalid_choice.is_some() {
return Err(Error::new(
proc_macro2::Span::call_site(),
"Enum choices must match the field type",
));
}
}
Ok(attributes)
}
}
fn choices_to_json_values(choices: Vec<Choice>) -> Result<Vec<serde_json::Value>> {
choices
.into_iter()
.map(|choice| choice.to_json_value())
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn build_parser(input: DeriveInput) -> OutputParser {
let mut parser = OutputParser::default();
parser.output_data.name = input.ident.to_string();
parser.output_data.schema._type = JsonType::Object.to_string();
parser.parse_struct_attributes(&input.attrs).unwrap();
parser.parse_data(input.data).unwrap();
parser
}
#[test]
fn parse_struct_populates_schema_and_required_fields() {
let input: DeriveInput = syn::parse_str(
r#"
#[strict(true)]
/// Greeting output
struct MyOutput {
#[output(description = "Name")]
name: String,
#[output(description = "Age")]
age: Option<u32>,
#[output(description = "Mode", choice = ["fast", "slow"])]
mode: String,
}
"#,
)
.unwrap();
let parser = build_parser(input);
assert_eq!(
parser.output_data.description.as_deref(),
Some("Greeting output")
);
assert_eq!(parser.output_data.strict, Some(true));
assert!(
parser
.output_data
.schema
.required
.contains(&"name".to_string())
);
assert!(
!parser
.output_data
.schema
.required
.contains(&"age".to_string())
);
let mode = parser.output_data.schema.properties.get("mode").unwrap();
assert_eq!(mode._type, "string");
assert_eq!(mode._enum.as_ref().unwrap().len(), 2);
assert_eq!(
mode._enum.as_ref().unwrap()[0],
serde_json::Value::String("fast".to_string())
);
let age = parser.output_data.schema.properties.get("age").unwrap();
assert_eq!(age._type, "integer");
}
#[test]
fn numeric_choices_serialize_as_json_numbers() {
let input: DeriveInput = syn::parse_str(
r#"
struct MyOutput {
#[output(description = "Level", choice = [1, 2, 3])]
level: u32,
}
"#,
)
.unwrap();
let parser = build_parser(input);
let level = parser.output_data.schema.properties.get("level").unwrap();
assert_eq!(level._type, "integer");
assert_eq!(
level._enum.as_ref().unwrap(),
&[
serde_json::json!(1),
serde_json::json!(2),
serde_json::json!(3),
]
);
}
#[test]
fn u64_choices_serialize_as_json_numbers() {
let input: DeriveInput = syn::parse_str(
r#"
struct MyOutput {
#[output(description = "Large id", choice = [10000000000000000000])]
id: u64,
}
"#,
)
.unwrap();
let parser = build_parser(input);
let id = parser.output_data.schema.properties.get("id").unwrap();
assert_eq!(id._type, "integer");
assert_eq!(
id._enum.as_ref().unwrap(),
&[serde_json::json!(10000000000000000000_u64)]
);
}
#[test]
fn multiline_doc_comments_join_description() {
let input: DeriveInput = syn::parse_str(
r#"
/// First line
/// Second line
struct MyOutput {
#[output(description = "Value")]
value: i64,
}
"#,
)
.unwrap();
let parser = build_parser(input);
assert_eq!(
parser.output_data.description.as_deref(),
Some("First line Second line")
);
}
#[test]
fn missing_output_attribute_errors() {
let input: DeriveInput = syn::parse_str(
r#"
struct BadOutput {
name: String,
}
"#,
)
.unwrap();
let mut parser = OutputParser::default();
parser.output_data.name = input.ident.to_string();
parser.output_data.schema._type = JsonType::Object.to_string();
let err = parser.parse_data(input.data).unwrap_err();
assert!(
err.to_string()
.contains("AgentOutput structs must have at least one field")
);
}
#[test]
fn tuple_struct_errors() {
let input: DeriveInput = syn::parse_str(r#"struct BadOutput(u32);"#).unwrap();
let mut parser = OutputParser::default();
parser.output_data.name = input.ident.to_string();
parser.output_data.schema._type = JsonType::Object.to_string();
let err = parser.parse_data(input.data).unwrap_err();
assert!(
err.to_string()
.contains("Tuple or Unit structs not yet supported")
);
}
}