klieo-macros 0.40.0

//! Procedural macros for the klieo agent framework.
//!
//! See the docs on [`tool`] for the `#[tool]` attribute macro and on
//! [`derive_klieo_response`] for the `#[derive(KlieoResponse)]` macro.

extern crate proc_macro;

mod codegen;
mod parse;

use parse::{ToolAttr, ToolDecl, ToolFn};
use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, spanned::Spanned, DeriveInput, FnArg, ItemFn, Pat, Type};

/// Turn an `async fn` into a `klieo_core::Tool` implementation.
///
/// # Example
///
/// ```ignore
/// use klieo_core::tool::ToolCtx;
/// use klieo_core::error::ToolError;
///
/// #[klieo_macros::tool(description = "Greet a person")]
/// async fn greet(_ctx: &ToolCtx, name: String) -> Result<String, ToolError> {
///     Ok(format!("Hello, {name}"))
/// }
/// ```
///
/// Generates a `GreetTool` zero-sized struct that implements
/// `klieo_core::Tool`. The user crate must declare:
/// `klieo-core`, `serde`, `serde_json`, `schemars`, `async-trait`.
#[proc_macro_attribute]
pub fn tool(attr: TokenStream, item: TokenStream) -> TokenStream {
    let attr = parse_macro_input!(attr as ToolAttr);
    let item_fn = parse_macro_input!(item as ItemFn);
    let func = match ToolFn::parse(item_fn) {
        Ok(f) => f,
        Err(e) => return e.to_compile_error().into(),
    };
    let decl = ToolDecl { attr, func };
    codegen::expand(decl).into()
}

/// Mark an `async fn` test as a klieo test: emits a sync
/// `#[test] fn` that builds a `current_thread` tokio runtime, optionally
/// constructs a `klieo_core::test_utils::TestContext`, and `block_on`s
/// the body.
///
/// # Accepted signatures
///
/// ```ignore
/// #[klieo::test]
/// async fn name() { /* … */ }
///
/// #[klieo::test]
/// async fn name(ctx: TestContext) { /* … */ }
/// ```
///
/// In the second form the macro builds `TestContext::default()` and
/// passes it in by value, so the test body can chain builder calls
/// (`ctx.with_canned_llm_responses(...)`) without touching `Arc<dyn …>`
/// plumbing.
///
/// The emitted code uses absolute `::tokio` / `::klieo_core` paths,
/// so callers do not need either crate in scope.
#[proc_macro_attribute]
pub fn test(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let item_fn = parse_macro_input!(item as ItemFn);

    if item_fn.sig.asyncness.is_none() {
        return syn::Error::new_spanned(
            item_fn.sig.fn_token,
            "#[klieo::test] requires an async fn",
        )
        .to_compile_error()
        .into();
    }

    let fn_name = &item_fn.sig.ident;
    let body = &item_fn.block;
    let attrs = &item_fn.attrs;
    let vis = &item_fn.vis;
    let inputs = &item_fn.sig.inputs;

    let injects_ctx = match inputs.len() {
        0 => false,
        1 => {
            // Must be a typed pattern (no self), an ident pattern, and
            // the type tail must be `TestContext`. We accept any module
            // prefix because absolute paths (`::klieo_core::…`) and
            // unqualified imports (`TestContext`) are both common in
            // tests.
            let arg = inputs.first().unwrap();
            let pat_ty = match arg {
                FnArg::Typed(pt) => pt,
                FnArg::Receiver(_) => {
                    return syn::Error::new_spanned(arg, "#[klieo::test] cannot decorate methods")
                        .to_compile_error()
                        .into();
                }
            };
            if !matches!(*pat_ty.pat, Pat::Ident(_)) {
                return syn::Error::new_spanned(
                    &pat_ty.pat,
                    "#[klieo::test] requires the argument to be `name: TestContext`",
                )
                .to_compile_error()
                .into();
            }
            if !is_test_context_type(&pat_ty.ty) {
                return syn::Error::new_spanned(
                    &pat_ty.ty,
                    "#[klieo::test] argument must be `TestContext` (the type-name tail must be `TestContext`)",
                )
                .to_compile_error()
                .into();
            }
            true
        }
        _ => {
            return syn::Error::new(
                inputs.span(),
                "#[klieo::test] accepts at most one argument: `ctx: TestContext`",
            )
            .to_compile_error()
            .into();
        }
    };

    let runtime_build = quote! {
        ::tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("tokio current-thread runtime builds")
    };

    let expanded = if injects_ctx {
        // Re-emit the argument name (`ctx` in the canonical shape) so
        // the body's references to it resolve. We bind it to a fresh
        // `TestContext::default()`. We preserve the user-written type
        // annotation (`TestContext` or `klieo_core::…::TestContext`)
        // so the import the user added at the call site stays
        // referenced — otherwise `#[allow(unused_imports)]` would be
        // required for every test fn.
        let arg = inputs.first().unwrap();
        let pat_ty = match arg {
            FnArg::Typed(pt) => pt,
            _ => unreachable!("guarded above"),
        };
        let arg_pat = &pat_ty.pat;
        let arg_ty = &pat_ty.ty;
        quote! {
            #( #attrs )*
            #[::core::prelude::v1::test]
            #vis fn #fn_name() {
                #runtime_build .block_on(async move {
                    #[allow(unused_mut)]
                    let mut #arg_pat: #arg_ty =
                        ::klieo::__private::klieo_core::test_utils::TestContext::default();
                    #body
                });
            }
        }
    } else {
        quote! {
            #( #attrs )*
            #[::core::prelude::v1::test]
            #vis fn #fn_name() {
                #runtime_build .block_on(async move {
                    #body
                });
            }
        }
    };

    expanded.into()
}

/// `true` if `ty`'s final path segment is `TestContext`. Permissive on
/// purpose: callers may use the absolute path
/// `klieo_core::test_utils::TestContext`, a use-imported `TestContext`,
/// or an alias — we only assert the tail.
fn is_test_context_type(ty: &Type) -> bool {
    let Type::Path(tp) = ty else {
        return false;
    };
    tp.path
        .segments
        .last()
        .map(|s| s.ident == "TestContext")
        .unwrap_or(false)
}

/// Derive `klieo_core::KlieoResponse` for a struct so it can be the
/// typed payload of a structured LLM response.
///
/// The derive only emits the `KlieoResponse` impl itself. The
/// generated `json_schema()` body delegates to
/// `schemars::schema_for!(Self)`, so the caller MUST also derive
/// `schemars::JsonSchema` and `serde::Deserialize` on the same type.
/// Keeping the schema and deserialize derives explicit makes the
/// dependency surface visible at the call site rather than hidden
/// inside this macro's expansion.
///
/// # Required derives
///
/// ```ignore
/// use klieo_macros::KlieoResponse;
/// use schemars::JsonSchema;
/// use serde::Deserialize;
///
/// #[derive(Deserialize, JsonSchema, KlieoResponse)]
/// struct Greeting {
///     greeting: String,
/// }
/// ```
///
/// The user crate must declare: `klieo-core`, `serde`, `serde_json`,
/// `schemars`.
#[proc_macro_derive(KlieoResponse)]
pub fn derive_klieo_response(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = &input.ident;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    let expanded = quote! {
        impl #impl_generics ::klieo::__private::klieo_core::response::KlieoResponse for #name #ty_generics
        #where_clause
        {
            fn json_schema() -> ::serde_json::Value {
                let schema = ::schemars::schema_for!(#name);
                ::serde_json::to_value(&schema)
                    .expect("schemars produces a valid serde_json::Value")
            }
        }
    };
    expanded.into()
}