llm_toolkit/

prompt.rs

//! A trait and macros for powerful, type-safe prompt generation.

use minijinja::Environment;
use serde::Serialize;

/// Represents a part of a multimodal prompt.
///
/// This enum allows prompts to contain different types of content,
/// such as text and images, enabling multimodal LLM interactions.
#[derive(Debug, Clone)]
pub enum PromptPart {
    /// Text content in the prompt.
    Text(String),
    /// Image content with media type and binary data.
    Image {
        /// The MIME media type (e.g., "image/jpeg", "image/png").
        media_type: String,
        /// The raw image data.
        data: Vec<u8>,
    },
    // Future variants like Audio or Video can be added here
}

/// A trait for converting any type into a string suitable for an LLM prompt.
///
/// This trait provides a standard interface for converting various types
/// into strings that can be used as prompts for language models.
///
/// # Example
///
/// ```
/// use llm_toolkit::prompt::ToPrompt;
///
/// // Common types have ToPrompt implementations
/// let number = 42;
/// assert_eq!(number.to_prompt(), "42");
///
/// let text = "Hello, LLM!";
/// assert_eq!(text.to_prompt(), "Hello, LLM!");
/// ```
///
/// # Custom Implementation
///
/// You can also implement `ToPrompt` directly for your own types:
///
/// ```
/// use llm_toolkit::prompt::{ToPrompt, PromptPart};
/// use std::fmt;
///
/// struct CustomType {
///     value: String,
/// }
///
/// impl fmt::Display for CustomType {
///     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
///         write!(f, "{}", self.value)
///     }
/// }
///
/// // By implementing ToPrompt directly, you can control the conversion.
/// impl ToPrompt for CustomType {
///     fn to_prompt_parts(&self) -> Vec<PromptPart> {
///         vec![PromptPart::Text(self.to_string())]
///     }
///
///     fn to_prompt(&self) -> String {
///         self.to_string()
///     }
/// }
///
/// let custom = CustomType { value: "custom".to_string() };
/// assert_eq!(custom.to_prompt(), "custom");
/// ```
pub trait ToPrompt {
    /// Converts the object into a vector of `PromptPart`s based on a mode.
    ///
    /// This is the core method that `derive(ToPrompt)` will implement.
    /// The `mode` argument allows for different prompt representations, such as:
    /// - "full": A comprehensive prompt with schema and examples.
    /// - "schema_only": Just the data structure's schema.
    /// - "example_only": Just a concrete example.
    ///
    /// The default implementation ignores the mode and calls `to_prompt_parts`
    /// for backward compatibility with manual implementations.
    fn to_prompt_parts_with_mode(&self, mode: &str) -> Vec<PromptPart> {
        // Default implementation for backward compatibility
        let _ = mode; // Unused in default impl
        self.to_prompt_parts()
    }

    /// Converts the object into a prompt string based on a mode.
    ///
    /// This method extracts only the text portions from `to_prompt_parts_with_mode()`.
    fn to_prompt_with_mode(&self, mode: &str) -> String {
        self.to_prompt_parts_with_mode(mode)
            .iter()
            .filter_map(|part| match part {
                PromptPart::Text(text) => Some(text.as_str()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join("")
    }

    /// Converts the object into a vector of `PromptPart`s using the default "full" mode.
    ///
    /// This method enables multimodal prompt generation by returning
    /// a collection of prompt parts that can include text, images, and
    /// other media types.
    fn to_prompt_parts(&self) -> Vec<PromptPart> {
        self.to_prompt_parts_with_mode("full")
    }

    /// Converts the object into a prompt string using the default "full" mode.
    ///
    /// This method provides backward compatibility by extracting only
    /// the text portions from `to_prompt_parts()` and joining them.
    fn to_prompt(&self) -> String {
        self.to_prompt_with_mode("full")
    }
}

// Add implementations for common types

impl ToPrompt for String {
    fn to_prompt_parts(&self) -> Vec<PromptPart> {
        vec![PromptPart::Text(self.clone())]
    }

    fn to_prompt(&self) -> String {
        self.clone()
    }
}

impl ToPrompt for &str {
    fn to_prompt_parts(&self) -> Vec<PromptPart> {
        vec![PromptPart::Text(self.to_string())]
    }

    fn to_prompt(&self) -> String {
        self.to_string()
    }
}

impl ToPrompt for bool {
    fn to_prompt_parts(&self) -> Vec<PromptPart> {
        vec![PromptPart::Text(self.to_string())]
    }

    fn to_prompt(&self) -> String {
        self.to_string()
    }
}

impl ToPrompt for char {
    fn to_prompt_parts(&self) -> Vec<PromptPart> {
        vec![PromptPart::Text(self.to_string())]
    }

    fn to_prompt(&self) -> String {
        self.to_string()
    }
}

macro_rules! impl_to_prompt_for_numbers {
    ($($t:ty),*) => {
        $(
            impl ToPrompt for $t {
                fn to_prompt_parts(&self) -> Vec<PromptPart> {
                    vec![PromptPart::Text(self.to_string())]
                }

                fn to_prompt(&self) -> String {
                    self.to_string()
                }
            }
        )*
    };
}

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

/// Renders a prompt from a template string and a serializable context.
///
/// This is the underlying function for the `prompt!` macro.
pub fn render_prompt<T: Serialize>(template: &str, context: T) -> Result<String, minijinja::Error> {
    let mut env = Environment::new();
    env.add_template("prompt", template)?;
    let tmpl = env.get_template("prompt")?;
    tmpl.render(context)
}

/// Creates a prompt string from a template and key-value pairs.
///
/// This macro provides a `println!`-like experience for building prompts
/// from various data sources. It leverages `minijinja` for templating.
///
/// # Example
///
/// ```
/// use llm_toolkit::prompt;
/// use serde::Serialize;
///
/// #[derive(Serialize)]
/// struct User {
///     name: &'static str,
///     role: &'static str,
/// }
///
/// let user = User { name: "Mai", role: "UX Engineer" };
/// let task = "designing a new macro";
///
/// let p = prompt!(
///     "User {{user.name}} ({{user.role}}) is currently {{task}}.",
///     user = user,
///     task = task
/// ).unwrap();
///
/// assert_eq!(p, "User Mai (UX Engineer) is currently designing a new macro.");
/// ```
#[macro_export]
macro_rules! prompt {
    ($template:expr, $($key:ident = $value:expr),* $(,)?) => {
        $crate::prompt::render_prompt($template, minijinja::context!($($key => $value),*))
    };
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde::Serialize;
    use std::fmt::Display;

    enum TestEnum {
        VariantA,
        VariantB,
    }

    impl Display for TestEnum {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            match self {
                TestEnum::VariantA => write!(f, "Variant A"),
                TestEnum::VariantB => write!(f, "Variant B"),
            }
        }
    }

    impl ToPrompt for TestEnum {
        fn to_prompt_parts(&self) -> Vec<PromptPart> {
            vec![PromptPart::Text(self.to_string())]
        }

        fn to_prompt(&self) -> String {
            self.to_string()
        }
    }

    #[test]
    fn test_to_prompt_for_enum() {
        let variant = TestEnum::VariantA;
        assert_eq!(variant.to_prompt(), "Variant A");
    }

    #[test]
    fn test_to_prompt_for_enum_variant_b() {
        let variant = TestEnum::VariantB;
        assert_eq!(variant.to_prompt(), "Variant B");
    }

    #[test]
    fn test_to_prompt_for_string() {
        let s = "hello world";
        assert_eq!(s.to_prompt(), "hello world");
    }

    #[test]
    fn test_to_prompt_for_number() {
        let n = 42;
        assert_eq!(n.to_prompt(), "42");
    }

    #[derive(Serialize)]
    struct SystemInfo {
        version: &'static str,
        os: &'static str,
    }

    #[test]
    fn test_prompt_macro_simple() {
        let user = "Yui";
        let task = "implementation";
        let prompt = prompt!(
            "User {{user}} is working on the {{task}}.",
            user = user,
            task = task
        )
        .unwrap();
        assert_eq!(prompt, "User Yui is working on the implementation.");
    }

    #[test]
    fn test_prompt_macro_with_struct() {
        let sys = SystemInfo {
            version: "0.1.0",
            os: "Rust",
        };
        let prompt = prompt!("System: {{sys.version}} on {{sys.os}}", sys = sys).unwrap();
        assert_eq!(prompt, "System: 0.1.0 on Rust");
    }

    #[test]
    fn test_prompt_macro_mixed() {
        let user = "Mai";
        let sys = SystemInfo {
            version: "0.1.0",
            os: "Rust",
        };
        let prompt = prompt!(
            "User {{user}} is using {{sys.os}} v{{sys.version}}.",
            user = user,
            sys = sys
        )
        .unwrap();
        assert_eq!(prompt, "User Mai is using Rust v0.1.0.");
    }

    #[test]
    fn test_prompt_macro_no_args() {
        let prompt = prompt!("This is a static prompt.",).unwrap();
        assert_eq!(prompt, "This is a static prompt.");
    }
}

#[derive(Debug, thiserror::Error)]
pub enum PromptSetError {
    #[error("Target '{target}' not found. Available targets: {available:?}")]
    TargetNotFound {
        target: String,
        available: Vec<String>,
    },
    #[error("Failed to render prompt for target '{target}': {source}")]
    RenderFailed {
        target: String,
        source: minijinja::Error,
    },
}

/// A trait for types that can generate multiple named prompt targets.
///
/// This trait enables a single data structure to produce different prompt formats
/// for various use cases (e.g., human-readable vs. machine-parsable formats).
///
/// # Example
///
/// ```ignore
/// use llm_toolkit::prompt::{ToPromptSet, PromptPart};
/// use serde::Serialize;
///
/// #[derive(ToPromptSet, Serialize)]
/// #[prompt_for(name = "Visual", template = "## {{title}}\n\n> {{description}}")]
/// struct Task {
///     title: String,
///     description: String,
///
///     #[prompt_for(name = "Agent")]
///     priority: u8,
///
///     #[prompt_for(name = "Agent", rename = "internal_id")]
///     id: u64,
///
///     #[prompt_for(skip)]
///     is_dirty: bool,
/// }
///
/// let task = Task {
///     title: "Implement feature".to_string(),
///     description: "Add new functionality".to_string(),
///     priority: 1,
///     id: 42,
///     is_dirty: false,
/// };
///
/// // Generate visual prompt
/// let visual_prompt = task.to_prompt_for("Visual")?;
///
/// // Generate agent prompt
/// let agent_prompt = task.to_prompt_for("Agent")?;
/// ```
pub trait ToPromptSet {
    /// Generates multimodal prompt parts for the specified target.
    fn to_prompt_parts_for(&self, target: &str) -> Result<Vec<PromptPart>, PromptSetError>;

    /// Generates a text prompt for the specified target.
    ///
    /// This method extracts only the text portions from `to_prompt_parts_for()`
    /// and joins them together.
    fn to_prompt_for(&self, target: &str) -> Result<String, PromptSetError> {
        let parts = self.to_prompt_parts_for(target)?;
        let text = parts
            .iter()
            .filter_map(|part| match part {
                PromptPart::Text(text) => Some(text.as_str()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join("\n");
        Ok(text)
    }
}

/// A trait for generating a prompt for a specific target type.
///
/// This allows a type (e.g., a `Tool`) to define how it should be represented
/// in a prompt when provided with a target context (e.g., an `Agent`).
pub trait ToPromptFor<T> {
    /// Generates a prompt for the given target, using a specific mode.
    fn to_prompt_for_with_mode(&self, target: &T, mode: &str) -> String;

    /// Generates a prompt for the given target using the default "full" mode.
    ///
    /// This method provides backward compatibility by calling the `_with_mode`
    /// variant with a default mode.
    fn to_prompt_for(&self, target: &T) -> String {
        self.to_prompt_for_with_mode(target, "full")
    }
}