aimx/inference/

response.rs

//! Inference response parsing and conversion
//!
//! This module provides functionality for parsing inference model responses and converting them
//! into structured data according to type definitions. It handles both single-line and multi-line
//! responses with various formatting patterns.
//!
//! # Overview
//!
//! The response parser processes text output from inference models and extracts structured
//! key-value pairs according to AIM response format conventions. It supports:
//! - Single-line responses with inline values
//! - Multi-line list responses with common prefixes
//! - Task status indicators (checkboxes)
//! - Type-safe conversion to [`Value`] and [`Literal`] types
//!
//! # Response Format
//!
//! The parser expects responses in the following format:
//! ```text
//! KEY1: inline value
//! KEY2
//! - Checkbox item 1 [X]
//! - Checkbox item 2 [ ]
//! KEY3
//! 1. Ordered item 1
//! 2. Ordered item 2
//! ```
//!
//! # Examples
//!
//! ```rust
//! use aimx::{parse_response, inference::Response};
//! use aimx::typedef::Typedef;
//! 
//! let response_text = "ANSWER: 42";
//! if let Some(response_map) = parse_response(response_text) {
//!     let typedef = Typedef::Number;
//!     if let Some(response) = response_map.get("ANSWER") {
//!         match response.to_value(&typedef) {
//!             Ok(value) => println!("Parsed value: {:?}", value),
//!             Err(e) => eprintln!("Error: {}", e),
//!         }
//!     }
//! }
//! ```

use std::collections::HashMap;
use crate::{
    Literal,
    inference::{Item, parse_key, parse_item, parse_inline_item, Suffix},
    literals::{parse_date, parse_number},
    Value,
    Typedef,
};
use anyhow::{anyhow, Result};

/// Represents a parsed inference response
/// 
/// This enum can represent either a single response item or a multi-line list of items.
/// The parser determines the response type based on the formatting pattern found in the input.
#[derive(Debug, PartialEq, Clone)]
pub enum Response {
    /// Single response item, typically from an inline key-value pair like "KEY: value"
    Single(Item),
    /// Multi-line response containing a list of items with common prefix patterns
    Multiline(Vec<Item>),
}

/// Parse a formatted inference response and extract key/item data
/// 
/// This function processes text output from inference models and extracts structured
/// key-value pairs according to AIM response format conventions. It handles both
/// single-line responses (with inline values) and multi-line list responses.
/// 
/// # Parameters
/// 
/// - `input`: The raw text response from an inference model
/// 
/// # Returns
/// 
/// - `Some(HashMap<String, Response>)`: A map of response keys to parsed responses
/// - `None`: If no valid responses were found in the input
/// 
/// # Examples
/// 
/// ```rust
/// use aimx::parse_response;
/// 
/// let response = "ANSWER: 42\nLIST:\n- Item 1\n- Item 2";
/// if let Some(parsed) = parse_response(response) {
///     assert!(parsed.contains_key("ANSWER"));
///     assert!(parsed.contains_key("LIST"));
/// }
/// ```
pub fn parse_response(input: &str) -> Option<HashMap<String, Response>> {
    let mut response_map: HashMap<String, Response> = HashMap::new();
    let mut current_list: Vec<Item> = Vec::new(); 
    let mut current_key: Option<String> = None;
    let mut first_item: Option<Item> = None;
    
    // Process each line
    for line in input.lines() {

        // Check if the line is a new key
        if let Ok((input, (key, suffix))) = parse_key(line) {
            // This marks a new response so insert any pending response
            if let (Some(key), Some(item)) = (current_key, first_item) {
                response_map.insert(key, Response::new(&item, &current_list));
            }
            match suffix {
                Suffix::Colon => {
                    // Store key and try to parse inline item
                    if let Ok((_, item)) = parse_inline_item(input) {
                        response_map.insert(key, Response::Single(item));
                    }
                    current_key = None;
                },
                Suffix::Eol | Suffix::ColonEol => {
                    // Store the new key which changes the state to multiline
                    current_key = Some(key);
                }
            }
            first_item = None;
            current_list.clear();

        // Check for multiline state
        } else if let Some(key) = &current_key {
            // Check if we have a first item
            if let Some(item) = &first_item {
                if let Ok((_, result)) = parse_item(line) {
                    // Check the prefixes match
                    let prefixes_match = match (&item, &result) {
                        (Item::Value(prefix1, _), Item::Value(prefix2, _)) => prefix1 == prefix2,
                        (Item::Task(prefix1, _, _), Item::Task(prefix2, _, _)) => prefix1 == prefix2,
                        _ => false, // Different item types
                    };
                    if prefixes_match {
                        current_list.push(result);
                    }
                    else {
                        response_map.insert(key.clone(), Response::new(&item, &current_list));
                        current_key = None;
                        first_item = None;
                        current_list.clear();
                    }
                }
                else {
                    response_map.insert(key.clone(), Response::new(&item, &current_list));
                    current_key = None;
                    first_item = None;
                    current_list.clear();
                }
            }
            // Parse the first item
            else if let Ok((_, result)) = parse_item(line) {
                current_list.push(result.clone());
                first_item = Some(result);
            }
            // No result
            else {
                current_key = None;
            }
        }
    }
    // This marks the end of the input so insert any pending response
    if let (Some(key), Some(item)) = (current_key, first_item) {
        response_map.insert(key, Response::new(&item, &current_list));
    }
     
    // Only return the map if we successfully parsed at least one item
    if response_map.is_empty() {
        None
    } else {
        Some(response_map)
    }
}

/// Convert a single input string to a Literal based on type definition
/// 
/// This function performs type-safe parsing of single values according to the
/// provided type definition. It supports number, date, and text types.
/// 
/// # Parameters
/// 
/// - `input`: The input string to parse
/// - `typedef`: The type definition specifying expected data type
/// 
/// # Returns
/// 
/// - `Ok(Literal)`: Successfully parsed literal value
/// - `Err`: If parsing fails or type mismatch occurs
fn from_single(input: &str, typedef: &Typedef) -> Result<Literal> {
    if input.trim().is_empty() {
        return Err(anyhow!("Nothing to assign"));
    }
    if typedef.is_number() {
        match parse_number(input) {
            Ok((_, number)) => {
                return Ok(Literal::from_number(number))
            }
            Err(e) => {
                return Err(anyhow!("Parse error: {}", e))
            }
        }
    }
    if typedef.is_date() {
        match parse_date(input) {
            Ok((_, date)) => {
                return Ok(Literal::from_date(date))
            }
            Err(e) => {
                return Err(anyhow!("Parse error: {}", e))
            }
        }        
    }
    if typedef.is_text() {
        return Ok(Literal::from_text(input.to_string()))
    }
    Err(anyhow!("Type definition mismatch"))
}

/// Convert a task input string to a Literal based on type definition
/// 
/// This function handles task-specific parsing with status indicators.
/// If the type definition doesn't specify a task, it falls back to regular parsing.
/// 
/// # Parameters
/// 
/// - `input`: The input string to parse
/// - `typedef`: The type definition specifying expected data type
/// - `status`: Optional task completion status
/// 
/// # Returns
/// 
/// - `Ok(Literal)`: Successfully parsed literal value
/// - `Err`: If parsing fails or type mismatch occurs
fn from_task(input: &str, typedef: &Typedef, status: &Option<bool>) -> Result<Literal> {
    if input.trim().is_empty() {
        return Err(anyhow!("Nothing to assign"));
    }
    if typedef.is_task() {
        return Ok(Literal::from_task(*status, input.to_string()))
    }
    // fall back to single
    from_single(input, typedef)
}

/// Convert a multi-line response to a Value based on type definition
/// 
/// This function processes a list of items from a multi-line response.
/// It can return either an array (if typedef specifies array) or the first valid value.
/// 
/// # Parameters
/// 
/// - `multiline`: Vector of items from multi-line response
/// - `typedef`: The type definition specifying expected data type
/// 
/// # Returns
/// 
/// - `Ok(Value)`: Successfully converted value (array or single value)
/// - `Err`: If no valid items found or type mismatch occurs
fn from_multiline(multiline: &Vec<Item>, typedef: &Typedef) -> Result<Value> {
    if multiline.is_empty() {
        return Err(anyhow!("Nothing to assign"));
    }
    let mut array: Vec<Box<Value>> = Vec::new();
    // Iterate multiline response
    for item in multiline.iter() {
        // Parse each item
        let result = match item {
            Item::Value(_, input) => {
                from_single(input, typedef)
            }
            Item::Task(_, status, input) => {
                from_task(input, typedef, status)
            }
        };
        // Accept valid type safe results
        if let Ok(literal) = result {
            // Build the array
            if typedef.is_array() {
                array.push(Box::new(Value::Literal(literal)));
            // Or return first value
            } else {
                return Ok(Value::Literal(literal));
            }
        }
    }
    if array.len() > 0 {
        return Ok(Value::Array(array));
    }
    Err(anyhow!("Type definition mismatch"))
}

impl Response {
    /// Create a new Response instance from an item and list
    /// 
    /// This constructor determines whether to create a Single or Multiline response
    /// based on the length of the provided list.
    /// 
    /// # Parameters
    /// 
    /// - `item`: The first item in the response
    /// - `list`: The complete list of items (including the first item)
    /// 
    /// # Returns
    /// 
    /// - `Response::Single`: If the list contains only one item
    /// - `Response::Multiline`: If the list contains multiple items
    pub fn new(item: &Item, list: &Vec<Item>) -> Self {
        if list.len() > 1 {
            Response::Multiline(list.clone())
        } else {
            Response::Single(item.clone())
        }
    }

    /// Convert the response to a Value based on type definition
    /// 
    /// This method performs type-safe conversion of the response content
    /// according to the provided type definition. It handles both single
    /// and multi-line responses appropriately.
    /// 
    /// # Parameters
    /// 
    /// - `typedef`: The type definition specifying expected data type
    /// 
    /// # Returns
    /// 
    /// - `Ok(Value)`: Successfully converted value
    /// - `Err`: If conversion fails due to type mismatch or parsing errors
    /// 
    /// # Examples
    /// 
    /// ```rust
    /// use aimx::inference::{Response, Item};
    /// use aimx::typedef::Typedef;
    /// use aimx::{Value, Literal, Prefix};
    /// 
    /// let item = Item::Value(Prefix::None, "42".to_string());
    /// let response = Response::Single(item);
    /// let typedef = Typedef::Number;
    /// 
    /// match response.to_value(&typedef) {
    ///     Ok(Value::Literal(Literal::Number(n))) => assert_eq!(n, 42.0),
    ///     _ => panic!("Unexpected result"),
    /// }
    /// ```
    pub fn to_value(&self, typedef: &Typedef) -> Result<Value> {
        match self {
            Response::Single(item) => {
                let literal = match item {
                    Item::Value(_, input) => {
                        from_single(input, typedef)?
                    }
                    Item::Task(_, status, input) => {
                        from_task(input, typedef, status)?
                    }
                };
                // Convert to array if typedef is array
                if typedef.is_array() {
                    let mut array: Vec<Box<Value>> = Vec::new();
                    array.push(Box::new(Value::Literal(literal)));
                    return Ok(Value::Array(array));
                }
                Ok(Value::Literal(literal))
            }
            Response::Multiline(multiline) => {
                from_multiline(multiline, typedef)
            }
        }
    }
}