cognis-core 0.2.0

Core traits and types for the Cognis LLM framework
Documentation
//! Support types and macros for tool derivation.
//!
//! - [`ToolJsonSchema`] trait — for framework-level schema resolution
//! - [`impl_tool!`] macro — generates `BaseTool` impl from a struct that has `json_schema()`

use serde_json::{json, Value};
use std::collections::{BTreeMap, HashMap};

/// Trait for types that can produce a JSON Schema representation.
///
/// This trait bridges the standalone `json_schema()` inherent method
/// (generated by `#[derive(JsonSchema)]`) with the cognis framework.
///
/// For most types, you should use `#[derive(JsonSchema)]` from `cognis_macros`
/// and then use [`impl_tool!`] to generate the `BaseTool` impl.
pub trait ToolJsonSchema {
    fn json_schema() -> Value;
}

// ---------------------------------------------------------------------------
// Primitive ToolJsonSchema implementations
// ---------------------------------------------------------------------------

impl ToolJsonSchema for String {
    fn json_schema() -> Value {
        json!({"type": "string"})
    }
}

impl ToolJsonSchema for &str {
    fn json_schema() -> Value {
        json!({"type": "string"})
    }
}

impl ToolJsonSchema for f32 {
    fn json_schema() -> Value {
        json!({"type": "number"})
    }
}

impl ToolJsonSchema for f64 {
    fn json_schema() -> Value {
        json!({"type": "number"})
    }
}

macro_rules! impl_integer_schema {
    ($($t:ty),*) => {
        $(
            impl ToolJsonSchema for $t {
                fn json_schema() -> Value {
                    json!({"type": "integer"})
                }
            }
        )*
    };
}

impl_integer_schema!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);

impl ToolJsonSchema for bool {
    fn json_schema() -> Value {
        json!({"type": "boolean"})
    }
}

impl<T: ToolJsonSchema> ToolJsonSchema for Vec<T> {
    fn json_schema() -> Value {
        json!({
            "type": "array",
            "items": T::json_schema()
        })
    }
}

impl<T: ToolJsonSchema> ToolJsonSchema for Option<T> {
    fn json_schema() -> Value {
        T::json_schema()
    }
}

impl<V: ToolJsonSchema> ToolJsonSchema for HashMap<String, V> {
    fn json_schema() -> Value {
        json!({
            "type": "object",
            "additionalProperties": V::json_schema()
        })
    }
}

impl<V: ToolJsonSchema> ToolJsonSchema for BTreeMap<String, V> {
    fn json_schema() -> Value {
        json!({
            "type": "object",
            "additionalProperties": V::json_schema()
        })
    }
}

impl ToolJsonSchema for Value {
    fn json_schema() -> Value {
        json!({})
    }
}

/// Generate a `BaseTool` implementation for a struct that has a `json_schema()`
/// inherent method (typically from `#[derive(JsonSchema)]`) and an
/// `async fn execute(&self) -> Result<ToolOutput>` method.
///
/// # Usage
///
/// ```ignore
/// use cognis_macros::JsonSchema;
/// use cognis_core::impl_tool;
///
/// #[derive(JsonSchema, Serialize, Deserialize)]
/// struct Calculator {
///     /// First operand
///     a: f64,
///     /// Second operand
///     b: f64,
/// }
///
/// impl Calculator {
///     async fn execute(&self) -> cognis_core::error::Result<cognis_core::tools::ToolOutput> {
///         Ok(cognis_core::tools::ToolOutput::Content(
///             serde_json::json!(self.a + self.b),
///         ))
///     }
/// }
///
/// impl_tool!(Calculator, name = "calculator", description = "Add two numbers");
/// ```
#[macro_export]
macro_rules! impl_tool {
    ($ty:ty, name = $name:expr, description = $desc:expr) => {
        #[async_trait::async_trait]
        impl cognis_core::tools::BaseTool for $ty {
            fn name(&self) -> &str {
                $name
            }

            fn description(&self) -> &str {
                $desc
            }

            fn args_schema(&self) -> Option<serde_json::Value> {
                Some(<$ty>::json_schema())
            }

            async fn _run(
                &self,
                _input: cognis_core::tools::ToolInput,
            ) -> cognis_core::error::Result<cognis_core::tools::ToolOutput> {
                self.execute().await
            }
        }

        impl cognis_core::tools::ToolJsonSchema for $ty {
            fn json_schema() -> serde_json::Value {
                <$ty>::json_schema()
            }
        }
    };
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_string_schema() {
        assert_eq!(String::json_schema(), json!({"type": "string"}));
    }

    #[test]
    fn test_f64_schema() {
        assert_eq!(f64::json_schema(), json!({"type": "number"}));
    }

    #[test]
    fn test_i32_schema() {
        assert_eq!(i32::json_schema(), json!({"type": "integer"}));
    }

    #[test]
    fn test_bool_schema() {
        assert_eq!(bool::json_schema(), json!({"type": "boolean"}));
    }

    #[test]
    fn test_vec_schema() {
        assert_eq!(
            Vec::<String>::json_schema(),
            json!({"type": "array", "items": {"type": "string"}})
        );
    }

    #[test]
    fn test_option_schema() {
        assert_eq!(Option::<i32>::json_schema(), json!({"type": "integer"}));
    }

    #[test]
    fn test_hashmap_schema() {
        assert_eq!(
            HashMap::<String, f64>::json_schema(),
            json!({"type": "object", "additionalProperties": {"type": "number"}})
        );
    }

    #[test]
    fn test_value_schema() {
        assert_eq!(Value::json_schema(), json!({}));
    }
}