aimx/inference/

response.rs

//! Inference response parsing and conversion.
//!
//! Parses model responses into structured items keyed by label and converts them to
//! [`Value`]s using [`Typedef`]s. Supports inline key-value pairs and multiline blocks
//! with common prefixes and checkbox-style task items. Permissive by design so that
//! strict validation can be applied in later stages.

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

/// Parsed inference response for a single key.
///
/// Either a single item or a multiline block extracted from model output.
#[derive(Debug, PartialEq, Clone)]
pub enum Response {
    /// Single response item, e.g. from an inline `KEY: value`.
    Single(Item),
    /// Multiline response with items sharing a common prefix/pattern.
    Multiline(Vec<Item>),
}

/// Parse formatted model output into a map from key to [`Response`].
///
/// Returns `None` if no valid keys or items are found.
pub fn parse_response(input: &str) -> Option<HashMap<Arc<str>, Response>> {
    let mut response_map: HashMap<Arc<str>, Response> = HashMap::new();
    let mut current_list: Vec<Item> = Vec::new();
    let mut current_key: Option<Arc<str>> = 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)
    }
}

/// Parse a single value using [`Typedef`] shape hints.
///
/// Supports number, date, and text typedefs; fails for incompatible typedefs or
/// empty input.
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(Arc::from(input)));
    }
    Err(anyhow!("Type definition mismatch"))
}

/// Parse a task-style item using [`Typedef`] hints.
///
/// Uses task typedefs when available; otherwise falls back to [`from_single`].
/// Fails on empty input.
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, Arc::from(input)));
    }
    // fall back to single
    from_single(input, typedef)
}

/// Convert a multiline block into a [`Value`] using [`Typedef`] hints.
///
/// For array typedefs collects all successfully parsed items; otherwise returns the
/// first successfully parsed item.
fn from_multiline(multiline: &[Item], typedef: &Typedef) -> Result<Arc<Value>> {
    if multiline.is_empty() {
        return Err(anyhow!("Nothing to assign"));
    }
    let mut array: Vec<Arc<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(Arc::new(Value::Literal(literal)));
            // Or return first value
            } else {
                return Ok(Arc::new(Value::Literal(literal)));
            }
        }
    }
    if !array.is_empty() {
        return Ok(Arc::new(Value::Array(Arc::new(array))));
    }
    Err(anyhow!("Type definition mismatch"))
}

impl Response {
    /// Create a [`Response`] from the collected items for a key.
    ///
    /// Uses [`Response::Single`] when `list` has one item, [`Response::Multiline`] otherwise.
    pub fn new(item: &Item, list: &[Item]) -> Self {
        if list.len() > 1 {
            Response::Multiline(list.to_vec())
        } else {
            Response::Single(item.clone())
        }
    }

    /// Convert this response to a [`Value`] using the given [`Typedef`].
    ///
    /// Respects array/task typedefs; errors on type mismatch or invalid content.
    pub fn to_value(&self, typedef: &Typedef) -> Result<Arc<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 array: Vec<Arc<Value>> = vec![Arc::new(Value::Literal(literal))];
                    return Ok(Arc::new(Value::Array(Arc::new(array))));
                }
                Ok(Arc::new(Value::Literal(literal)))
            }
            Response::Multiline(multiline) => from_multiline(multiline, typedef),
        }
    }
}