1use crate::detail::ffi::{
9 FOUNDRY_LOCAL_PARAM_DO_SAMPLE, FOUNDRY_LOCAL_PARAM_EARLY_STOPPING,
10 FOUNDRY_LOCAL_PARAM_FREQUENCY_PENALTY, FOUNDRY_LOCAL_PARAM_MAX_OUTPUT_TOKENS,
11 FOUNDRY_LOCAL_PARAM_PRESENCE_PENALTY, FOUNDRY_LOCAL_PARAM_SEED,
12 FOUNDRY_LOCAL_PARAM_TEMPERATURE, FOUNDRY_LOCAL_PARAM_TOOL_CHOICE, FOUNDRY_LOCAL_PARAM_TOP_K,
13 FOUNDRY_LOCAL_PARAM_TOP_P,
14};
15use crate::item::Item;
16use crate::item_queue::ItemQueue;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
20pub enum ToolChoice {
21 #[default]
23 Auto,
24 None,
26 Required,
28}
29
30impl ToolChoice {
31 fn as_param(self) -> &'static str {
32 match self {
33 ToolChoice::Auto => "auto",
34 ToolChoice::None => "none",
35 ToolChoice::Required => "required",
36 }
37 }
38}
39
40#[derive(Debug, Clone, PartialEq, Default)]
44pub struct SearchOptions {
45 pub temperature: Option<f32>,
47 pub top_p: Option<f32>,
49 pub top_k: Option<i32>,
51 pub max_output_tokens: Option<i32>,
53 pub frequency_penalty: Option<f32>,
55 pub presence_penalty: Option<f32>,
57 pub seed: Option<i64>,
59 pub early_stopping: Option<bool>,
61 pub do_sample: Option<bool>,
63}
64
65#[derive(Debug, Clone, PartialEq, Default)]
74pub struct RequestOptions {
75 pub search: SearchOptions,
77 pub tool_choice: Option<ToolChoice>,
79 pub additional_options: Vec<(String, String)>,
82}
83
84impl RequestOptions {
85 pub(crate) fn to_pairs(&self) -> Vec<(String, String)> {
90 let mut pairs: Vec<(String, String)> = self.additional_options.clone();
91 let s = &self.search;
92 if let Some(v) = s.temperature {
93 pairs.push((FOUNDRY_LOCAL_PARAM_TEMPERATURE.to_string(), v.to_string()));
94 }
95 if let Some(v) = s.top_p {
96 pairs.push((FOUNDRY_LOCAL_PARAM_TOP_P.to_string(), v.to_string()));
97 }
98 if let Some(v) = s.top_k {
99 pairs.push((FOUNDRY_LOCAL_PARAM_TOP_K.to_string(), v.to_string()));
100 }
101 if let Some(v) = s.max_output_tokens {
102 pairs.push((
103 FOUNDRY_LOCAL_PARAM_MAX_OUTPUT_TOKENS.to_string(),
104 v.to_string(),
105 ));
106 }
107 if let Some(v) = s.frequency_penalty {
108 pairs.push((
109 FOUNDRY_LOCAL_PARAM_FREQUENCY_PENALTY.to_string(),
110 v.to_string(),
111 ));
112 }
113 if let Some(v) = s.presence_penalty {
114 pairs.push((
115 FOUNDRY_LOCAL_PARAM_PRESENCE_PENALTY.to_string(),
116 v.to_string(),
117 ));
118 }
119 if let Some(v) = s.seed {
120 pairs.push((FOUNDRY_LOCAL_PARAM_SEED.to_string(), v.to_string()));
121 }
122 if let Some(v) = s.early_stopping {
123 pairs.push((
124 FOUNDRY_LOCAL_PARAM_EARLY_STOPPING.to_string(),
125 bool_str(v).to_string(),
126 ));
127 }
128 if let Some(v) = s.do_sample {
129 pairs.push((
130 FOUNDRY_LOCAL_PARAM_DO_SAMPLE.to_string(),
131 bool_str(v).to_string(),
132 ));
133 }
134 if let Some(tc) = self.tool_choice {
135 pairs.push((
136 FOUNDRY_LOCAL_PARAM_TOOL_CHOICE.to_string(),
137 tc.as_param().to_string(),
138 ));
139 }
140 pairs
141 }
142}
143
144fn bool_str(v: bool) -> &'static str {
145 if v {
146 "true"
147 } else {
148 "false"
149 }
150}
151
152#[derive(Debug, Clone, Default)]
158pub struct Request {
159 pub items: Vec<Item>,
161 pub input_queue: Option<ItemQueue>,
164 pub options: Option<RequestOptions>,
166}
167
168impl Request {
169 pub fn new() -> Self {
171 Self::default()
172 }
173
174 pub fn from_items(items: impl Into<Vec<Item>>) -> Self {
176 Self {
177 items: items.into(),
178 input_queue: None,
179 options: None,
180 }
181 }
182
183 pub fn with_item(mut self, item: Item) -> Self {
185 self.items.push(item);
186 self
187 }
188
189 pub fn with_input_queue(mut self, queue: ItemQueue) -> Self {
191 self.input_queue = Some(queue);
192 self
193 }
194
195 pub fn with_options(mut self, options: RequestOptions) -> Self {
197 self.options = Some(options);
198 self
199 }
200
201 pub(crate) fn option_pairs(&self) -> Vec<(String, String)> {
203 self.options
204 .as_ref()
205 .map(RequestOptions::to_pairs)
206 .unwrap_or_default()
207 }
208}
209
210#[cfg(test)]
211mod tests {
212 use super::*;
213
214 #[test]
215 fn typed_fields_override_additional_options() {
216 let opts = RequestOptions {
217 search: SearchOptions {
218 temperature: Some(0.5),
219 max_output_tokens: Some(128),
220 do_sample: Some(false),
221 ..Default::default()
222 },
223 tool_choice: Some(ToolChoice::Required),
224 additional_options: vec![("temperature".into(), "9.9".into())],
225 };
226 let pairs = opts.to_pairs();
227 assert_eq!(pairs[0], ("temperature".to_string(), "9.9".to_string()));
229 assert!(pairs.iter().any(|(k, v)| k == "temperature" && v == "0.5"));
230 assert!(pairs.iter().rposition(|(k, _)| k == "temperature").unwrap() > 0);
231 assert!(pairs
232 .iter()
233 .any(|(k, v)| k == "max_output_tokens" && v == "128"));
234 assert!(pairs.iter().any(|(k, v)| k == "do_sample" && v == "false"));
235 assert!(pairs
236 .iter()
237 .any(|(k, v)| k == "tool_choice" && v == "required"));
238 }
239
240 #[test]
241 fn builder_assembles_request() {
242 let req = Request::from_items(vec![Item::text("hi")])
243 .with_item(Item::text("there"))
244 .with_options(RequestOptions::default());
245 assert_eq!(req.items.len(), 2);
246 assert!(req.options.is_some());
247 assert!(req.input_queue.is_none());
248 }
249}