aimx/inference/

key.rs

//! Inference key parsing for Agentic Inference Markup (AIM) format.
//!
//! This module provides functionality for parsing inference keys, which are UPPERCASE C-style identifiers
//! used in AIM files to identify different types of inference directives and their termination patterns.
//! These keys are used to structure inference prompts and parse model responses.
//!
//! # Key Grammar
//! The key grammar supports three different suffix patterns:
//! - `key_colon_eol`: UCID ':' whitespace? end-of-line (e.g., `INSTRUCTIONS:\\n`) - indicates multi-line content follows
//! - `key_eol`: UCID whitespace? end-of-line (e.g., `SYSTEM\\n`) - indicates simple value on same line
//! - `key_colon`: UCID ':' (e.g., `MODEL:`) - indicates inline value follows
//!
//! # Examples
//! In AIM files, inference keys are used to structure prompts:
//! ```text
//! // Modifier rules in AIM files
//! SYSTEM: You are a helpful assistant
//! INSTRUCTIONS:
//! Summarize the following text in 100 words or less.
//! Focus on key points and maintain clarity.
//! MODEL: gpt-4
//! ```
//!
//! In model responses, keys indicate how to parse the returned content:
//! ```text
//! ANSWER: 42
//! LIST:
//! - First item
//! - Second item
//! SUMMARY
//! This is a summary on the same line.
//! ```
//!
//! The parser recognizes these keys and their different termination patterns to properly
//! structure inference prompts and parse responses.

use nom::{
    error::{Error, ErrorKind},
    IResult, Parser,
    character::complete::{satisfy, char, multispace0},
    bytes::complete::{take_while},
    combinator::recognize,
    branch::alt,
    sequence::pair,
};

/// Represents the different suffix patterns that can follow an inference key.
///
/// Inference keys in AIM files can have different termination patterns that indicate
/// how the key's value should be parsed and interpreted. The parser tries patterns
/// in a specific order of precedence.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Suffix {
    /// Key followed by colon and end-of-line (e.g., `INSTRUCTIONS:\\n`)
    /// This indicates a multi-line value follows on subsequent lines.
    /// This pattern has the highest precedence in parsing.
    ColonEol,
    
    /// Key followed by optional whitespace and end-of-line (e.g., `SYSTEM\\n`)
    /// This indicates a simple value follows on the same line or in the response parser
    /// indicates that list items follow.
    /// This pattern has medium precedence in parsing.
    Eol,
    
    /// Key followed by colon (e.g., `MODEL:`)
    /// This indicates an inline value follows on the same line.
    /// This pattern has the lowest precedence in parsing.
    Colon,
}

/// Parse an UPPERCASE C-style identifier (UCID).
///
/// UCIDs follow the pattern: starting with an uppercase letter or underscore,
/// followed by any combination of uppercase letters, digits, or underscores.
///
/// # Grammar
/// ```text
/// UCID = [A-Z_] [A-Z0-9_]*
/// ```
///
/// # Examples
/// ```rust
/// use aimx::inference::key::parse_ucid;
///
/// assert_eq!(parse_ucid("INSTRUCTIONS"), Ok(("", "INSTRUCTIONS")));
/// assert_eq!(parse_ucid("MODEL_NAME"), Ok(("", "MODEL_NAME")));
/// assert_eq!(parse_ucid("_PRIVATE_KEY"), Ok(("", "_PRIVATE_KEY")));
/// ```
///
/// # Arguments
/// * `input` - The input string to parse
///
/// # Returns
/// Returns an `IResult` containing the remaining input and the parsed identifier.
pub fn parse_ucid(input: &str) -> IResult<&str, &str> {
    recognize(pair(
        satisfy(|c| (c >= 'A' && c <= 'Z') || c == '_'),
        take_while(|c: char| (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_')
    )).parse(input)
}

/// Check if we're at the end of a line (no more content).
///
/// This is a helper function that succeeds only when there's no more content
/// to parse, indicating the end of a line or input.
fn check_eol(input: &str) -> IResult<&str, ()> {
    if input.is_empty() {
        Ok((input, ()))
    } else {
        Err(nom::Err::Error(Error::new(input, ErrorKind::Fail)))
    }
}

/// Parse an inference key according to the AIM grammar.
///
/// This function parses inference keys and their suffix patterns, which indicate
/// how the key's value should be interpreted. The parser tries each suffix pattern
/// in order of specificity (most specific first).
///
/// # Grammar Patterns
/// 1. `key_colon_eol` = UCID ':' whitespace? EOL
/// 2. `key_eol` = UCID whitespace? EOL
/// 3. `key_colon` = UCID ':'
///
/// # Examples
/// ```rust
/// use aimx::inference::key::{parse_key, Suffix};
///
/// // ColonEol pattern
/// assert_eq!( 
///     parse_key("INSTRUCTIONS:\n"), 
///     Ok(("", ("INSTRUCTIONS".to_string(), Suffix::ColonEol)))
/// );
///
/// // Eol pattern
/// assert_eq!(
///     parse_key("SYSTEM\n"),
///     Ok(("", ("SYSTEM".to_string(), Suffix::Eol)))
/// );
///
/// // Colon pattern
/// assert_eq!(
///     parse_key("MODEL:"),
///     Ok(("", ("MODEL".to_string(), Suffix::ColonEol)))
/// );
/// ```
///
/// # Arguments
/// * `input` - The input string to parse
///
/// # Returns
/// Returns an `IResult` containing the remaining input and a tuple of:
/// - The parsed key as a `String`
/// - The detected `Suffix` pattern
pub fn parse_key(input: &str) -> IResult<&str, (String, Suffix)> {
    let (input, _) = multispace0(input)?;
    let (remaining, ucid) = parse_ucid(input)?;
    
    // Try to parse each possible suffix in order of specificity
    alt((
        // key_colon_eol = UCID ':' WS? EOL
        |input| {
            let (input, _) = char(':')(input)?;
            let (input, _) = multispace0(input)?;
            let (input, _) = check_eol(input)?;
            Ok((input, (ucid.to_string(), Suffix::ColonEol)))
        },
        // key_eol = UCID WS? EOL  
        |input| {
            let (input, _) = multispace0(input)?;
            let (input, _) = check_eol(input)?;
            Ok((input, (ucid.to_string(), Suffix::Eol)))
        },
        // key_colon = UCID ':'
        |input| {
            let (input, _) = char(':')(input)?;
            Ok((input, (ucid.to_string(), Suffix::Colon)))
        }
    )).parse(remaining)
}