aimx/inference/
inference.rs1use 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#[derive(Debug, Clone, PartialEq)]
16pub enum Model {
17 Fast,
20 Standard,
23 Thinking,
26 Extract,
29 Instruct,
32 Coder,
35}
36
37impl Model {
38 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#[derive(Debug, PartialEq, Clone)]
70pub enum Pattern {
71 Evaluate,
73 Inference,
75 Search,
77 Summarize,
79 Compose,
81 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#[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
194fn 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}