aimx/inference/
inference.rs

1use crate::{
2    aim::{WriterLike, Writer},
3    expressions::parse_identifier
4};
5use nom::{
6    IResult, Parser,
7    character::complete::multispace0,
8    sequence::{delimited, preceded},
9};
10use std::{
11    fmt,
12};
13
14/// Inference model selection for AIMX workflows.
15#[derive(Debug, Clone, PartialEq)]
16pub enum Model {
17    /// Fast model - optimized for speed and low latency
18    /// Suitable for simple classification, extraction, or quick responses
19    Fast,
20    /// Standard model - balanced performance for general use
21    /// Suitable for most inference tasks including summarization and basic reasoning
22    Standard,
23    /// Thinking model - optimized for complex reasoning and planning
24    /// Suitable for multi-step reasoning, complex problem solving, and creative tasks
25    Thinking,
26    /// Extraction model - optimized for a large context window
27    /// Suitable for extraction tasks that need the largest available context window
28    Extract,
29    /// Instruct model - optimized for closely following instructions
30    /// Suitable for model as a judge and compliance tasks
31    Instruct,
32    /// Coder model - optimized for programming tasks
33    /// Suitable for code and aimx formula generation
34    Coder,
35}
36
37impl Model {
38    /// Map a string identifier to a [`Model`].
39    /// Unrecognized values, including "standard", return [`Model::Standard`].
40    pub fn new(identifier: &str) -> Option<Self> {
41        match identifier {
42            "fast" => Some(Model::Fast),
43            "standard" => Some(Model::Standard),
44            "thinking" => Some(Model::Thinking),
45            "extract" => Some(Model::Extract),
46            "instruct" => Some(Model::Instruct),
47            "coder" => Some(Model::Coder),
48            _ => None,
49        }
50    }
51
52    pub fn as_str(&self) -> &'static str {
53        match self {
54            Model::Fast => "fast",
55            Model::Standard => "standard",
56            Model::Thinking => "thinking",
57            Model::Extract => "extract",
58            Model::Instruct => "instruct",
59            Model::Coder => "coder",
60        }
61    }
62
63    pub fn print(&self, writer: &mut Writer) {
64        writer.write_str(self.as_str());
65    }
66}
67
68/// Inference orchestration pattern.
69#[derive(Debug, PartialEq, Clone)]
70pub enum Pattern {
71    // Non-agentic workflow evaluation
72    Evaluate,
73    // General inference pattern
74    Inference,
75    // Search extraction pattern
76    Search,
77    // Map-Reduce pattern
78    Summarize,
79    // React composer pattern for accurate annotation based document editing (Crate markdown)
80    Compose,
81    // Debate pattern where two agents each make an assessment; a third judge scores resolution.
82    Debate,
83}
84
85impl Pattern {
86    pub fn new(identifier: &str) -> Option<Self> {
87        match identifier {
88            "evaluate" => Some(Pattern::Evaluate),
89            "inference" => Some(Pattern::Inference),
90            "search" => Some(Pattern::Search),
91            "summarize" => Some(Pattern::Summarize),
92            "compose" => Some(Pattern::Compose),
93            "debate" => Some(Pattern::Debate),
94            _ => None,
95        }
96    }
97    pub fn as_str(&self) -> &'static str {
98        match self {
99            Pattern::Evaluate => "evaluate",
100            Pattern::Inference => "inference",
101            Pattern::Search => "search",
102            Pattern::Summarize => "summarize",
103            Pattern::Compose => "compose",
104            Pattern::Debate => "debate",
105         }
106    }
107
108    pub fn print(&self, writer: &mut Writer) {
109        writer.write_str(self.as_str());
110    }
111}
112
113/// Inference configuration: orchestration pattern and model.
114#[derive(Debug, PartialEq, Clone)]
115pub struct Inference {
116    pattern: Pattern,
117    model: Option<Model>,
118}
119
120impl Default for Inference {
121    fn default() -> Self {
122        Inference {
123            pattern: Pattern::Evaluate,
124            model: None,
125        }
126    }
127}
128
129impl Inference {
130    pub fn new(pattern: Pattern, model: Option<Model>) -> Self {
131        Self {
132            pattern,
133            model,
134        }
135    }
136
137    pub fn pattern(&self) -> &Pattern {
138        &self.pattern
139    }
140
141    pub fn model(&self) -> &Model {
142        match &self.model {
143            Some(model) => model,
144            None => &Model::Standard,
145        }
146    }
147
148    pub fn print(&self, writer: &mut Writer) {
149        let mut add_space = false;
150        if self.pattern != Pattern::Evaluate {
151            self.pattern.print(writer);
152            add_space = true;
153        }
154        if let Some(model) = &self.model {
155            if add_space {
156                writer.write_char(' ');
157            }
158            model.print(writer);
159        }
160    }
161}
162
163impl WriterLike for Inference {
164    fn write(&self, writer: &mut Writer) {
165        self.print(writer);
166    }
167}
168
169impl fmt::Display for Inference {
170    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171        write!(f, "{}", self.to_stringized())
172    }
173}
174
175pub fn parse_inference(input: &str) -> IResult<&str, Option<Inference>> {
176    let (input, (opt_pattern, opt_model)) = delimited(
177        multispace0,
178        parse_options,
179        multispace0,
180    ).parse(input)?;
181
182    let pattern = match opt_pattern {
183        Some(pattern) => pattern,
184        None => Pattern::Evaluate,
185    };
186
187    if pattern == Pattern::Evaluate && opt_model.is_none() {
188        return Ok((input, None));
189    }
190
191    Ok((input, Some(Inference::new(pattern, opt_model))))
192}
193
194/// Parse inference options: pattern, model, @state and !version (source template)
195/// Returns (pattern_opt, model_opt, state_opt) tuple
196fn parse_options(input: &str) -> IResult<&str, (Option<Pattern>, Option<Model>)> {
197    let mut opt_pattern: Option<Pattern> = None;
198    let mut opt_model: Option<Model> = None;
199
200    let mut input = input;
201    while let Ok((remain, identifier)) = preceded(
202        multispace0, 
203        parse_identifier).parse(input) {
204
205        if let Some(pattern) = Pattern::new(identifier) {
206            opt_pattern = Some(pattern);
207        } else if let Some(model) = Model::new(identifier) {
208            opt_model = Some(model);
209        } else {
210            break;
211        }
212        input = remain;
213    }
214    Ok((input, (opt_pattern, opt_model)))
215}