aimx/inference/

modifier.rs

use crate::{
    aim::{ContextLike, Prefix, Rule, WorkflowLike},
    expressions::ExpressionLike,
    inference::OutputFormat,
    values::Value,
};
use std::{collections::HashMap, sync::Arc};

/// Collects workflow modifier rules into a prompt description.
///
/// Built from a `WorkflowLike` and `ContextLike`. Interprets rules as:
/// - `ROLE`: system role/persona.
/// - `INSTRUCTIONS`: global instructions.
/// - `OUTPUT`: explicit output-format specification or free text.
/// - `SYSTEM`: overrides default system prompt string.
/// - `content`: main user content.
/// - UPPERCASE rules: additional modifiers in [`Modifier::others`].
/// - Assignment (`name`) + matching UPPERCASE format rules: captured as
///   field specifications and examples.
#[derive(Debug, PartialEq, Clone)]
pub struct Modifier {
    /// Value from `ROLE`.
    role: Option<Arc<Value>>,
    /// Value from `INSTRUCTIONS`.
    instructions: Option<Arc<Value>>,
    /// Value from `OUTPUT`.
    output_format: Option<Arc<Value>>,
    /// System prompt (default or `SYSTEM`).
    system: Arc<str>,
    /// User content from `content`.
    user: Option<Arc<Value>>,
    /// Other UPPERCASE modifiers.
    others: Vec<(Arc<str>, Arc<Value>)>,
    /// Assignment rules keyed by uppercased identifier.
    assignments: HashMap<Arc<str>, Arc<Rule>>,
    /// Assignment formats and examples.
    examples: Vec<(Arc<str>, Arc<Value>)>,
}

impl Default for Modifier {
    fn default() -> Self {
        Self {
            role: None,
            instructions: None,
            output_format: None,
            system: Arc::from("You are a helpful assistant"),
            user: None,
            others: Vec::new(),
            assignments: HashMap::new(),
            examples: Vec::new(),
        }
    }
}

impl Modifier {
    pub fn new(workflow: Arc<dyn WorkflowLike>, context: &mut dyn ContextLike) -> Self {
        let mut modifier = Self::default();

        // ----- FIND ASSIGNMENTS -----
        // Iterate each rule
        for rule in workflow.iter_rules() {
            // Select assignment rules (rule identifiers with leading underscores)
            if rule.is_assignment() {
                // Convert _identifier (leading underscore) to IDENTIFIER (UPPERCASE)
                let ucid = rule.ucid();
                // Create a map of ucid inference rules
                modifier.assignments.insert(ucid, rule.clone());
            }
        }

        // Iterate each rule
        for rule in workflow.iter_rules() {
            let identifier = rule.identifier();
            // content := USER TEXT (user prompt)
            if &*identifier == "content" {
                if let Ok(value) = rule.evaluate(context) {
                    modifier.user = Some(value);
                }
            // Select modifiers (rule identifiers that are UPPERCASE)
            } else if rule.is_modifier() {
                // Role persona
                if &*identifier == "ROLE" {
                    if let Ok(value) = rule.evaluate(context) {
                        modifier.role = Some(value);
                    }
                // Inference instructions
                } else if &*identifier == "INSTRUCTIONS" {
                    if let Ok(value) = rule.evaluate(context) {
                        modifier.instructions = Some(value);
                    }
                // Output format instructions and extraction source examples
                } else if &*identifier == "OUTPUT" {
                    if let Ok(value) = rule.evaluate(context) {
                        modifier.output_format = Some(value);
                    }
                // System prompt
                } else if &*identifier == "SYSTEM" {
                    if let Ok(value) = rule.evaluate(context) {
                        modifier.system = value.to_pretty_str(Prefix::None);
                    }
                // Specific rule format, specification and examples
                } else if rule.is_format() && modifier.assignments.contains_key(&identifier) {
                    if let Ok(value) = rule.evaluate(context) {
                        modifier.examples.push((identifier, value));
                    }
                // Other general prompt modifiers
                } else if let Ok(value) = rule.evaluate(context) {
                    modifier.others.push((identifier, value));
                }
            }
        }

        modifier
    }

    pub fn system_prompt(&self) -> Arc<str> {
        self.system.clone()
    }

    /// Write the role section.
    pub fn role(&self, format: &mut OutputFormat) {
        if let Some(value) = &self.role {
            format.value(value);
        } else {
            format.text("You are a helpful assistant.");
        }
        format.section_end();
    }

    /// Write all other UPPERCASE modifiers as sections.
    pub fn others(&self, format: &mut OutputFormat) {
        if !self.others.is_empty() {
            for (identifier, value) in &self.others {
                format.section_start(identifier);
                format.unordered_value(value);
            }
            format.section_end();
        }
    }

    /// Write INSTRUCTIONS from explicit rule or derived from assignment formats.
    pub fn instructions(&self, format: &mut OutputFormat) {
        format.section_start("INSTRUCTIONS");
        if let Some(value) = self.instructions.clone() {
            format.value(&value);
        } else if !self.examples.is_empty() {
            format.text("Derive the following values from the content:");
        }
        for (identifier, value) in self.examples.clone() {
            if let Value::Format(f) = &*value {
                format.single_item(&identifier, &f.instruction());
            }
        }
        format.section_end();
    }

    /// Write OUTPUT FORMAT from explicit rule or assignment formats.
    pub fn output_format(&self, format: &mut OutputFormat) {
        format.section_start("OUTPUT FORMAT");
        if let Some(value) = self.output_format.clone() {
            if let Value::Format(f) = &*value {
                format.text(&f.instruction());
            } else {
                format.value(&value);
            }
        } else if !self.examples.is_empty() {
            format.text("Provide your answer in this exact format:");
        }
        for (identifier, value) in self.examples.clone() {
            if let Value::Format(f) = &*value {
                format.single_line_tagged(&identifier, &f.template());
            }
        }
        format.section_end();
    }

    /// Write EXAMPLES from `OUTPUT` and assignment formats, if present.
    pub fn examples(&self, format: &mut OutputFormat) {
        format.section_start("EXAMPLES");
        let mut index: usize = 0;
        let mut length: usize = 1;
        while index < length {
            if let Some(value) = self.output_format.clone() && let Value::Format(f) = &*value {
                length = f.examples().len();
                if index < length {
                    format.single_line_quoted("USER TEXT", &f.examples()[index]);
                }
            }
            for (identifier, value) in self.examples.clone() {
                if let Value::Format(f) = &*value {
                    length = f.examples().len();
                    if index < length {
                        format.single_line(&identifier, &f.examples()[index]);
                    }
                }
            }
            format.section_end();
            index += 1;
        }
    }

    /// Write USER TEXT section if `content` exists.
    pub fn user_text(&self, format: &mut OutputFormat) {
        if let Some(value) = &self.user {
            format.section_start("USER TEXT");
            format.value(value);
            format.section_end();
        }
    }
}