langchain_rust/language_models/
options.rs

1use futures::Future;
2use std::{pin::Pin, sync::Arc};
3use tokio::sync::Mutex;
4
5use crate::schemas::{FunctionCallBehavior, FunctionDefinition, ResponseFormat};
6
7#[derive(Clone)]
8pub struct CallOptions {
9    pub candidate_count: Option<usize>,
10    pub max_tokens: Option<u32>,
11    pub temperature: Option<f32>,
12    pub stop_words: Option<Vec<String>>,
13    pub streaming_func: Option<
14        Arc<
15            Mutex<dyn FnMut(String) -> Pin<Box<dyn Future<Output = Result<(), ()>> + Send>> + Send>,
16        >,
17    >,
18    pub top_k: Option<usize>,
19    pub top_p: Option<f32>,
20    pub seed: Option<usize>,
21    pub min_length: Option<usize>,
22    pub max_length: Option<usize>,
23    pub n: Option<usize>,
24    pub repetition_penalty: Option<f32>,
25    pub frequency_penalty: Option<f32>,
26    pub presence_penalty: Option<f32>,
27    pub functions: Option<Vec<FunctionDefinition>>,
28    pub function_call_behavior: Option<FunctionCallBehavior>,
29    pub response_format: Option<ResponseFormat>,
30    pub stream_usage: Option<bool>,
31}
32
33impl Default for CallOptions {
34    fn default() -> Self {
35        CallOptions::new()
36    }
37}
38impl CallOptions {
39    pub fn new() -> Self {
40        CallOptions {
41            candidate_count: None,
42            max_tokens: None,
43            temperature: None,
44            stop_words: None,
45            streaming_func: None,
46            top_k: None,
47            top_p: None,
48            seed: None,
49            min_length: None,
50            max_length: None,
51            n: None,
52            repetition_penalty: None,
53            frequency_penalty: None,
54            presence_penalty: None,
55            functions: None,
56            function_call_behavior: None,
57            response_format: None,
58            stream_usage: None,
59        }
60    }
61
62    // Refactored "with" functions as methods of CallOptions
63    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
64        self.max_tokens = Some(max_tokens);
65        self
66    }
67
68    pub fn with_candidate_count(mut self, candidate_count: usize) -> Self {
69        self.candidate_count = Some(candidate_count);
70        self
71    }
72
73    pub fn with_temperature(mut self, temperature: f32) -> Self {
74        self.temperature = Some(temperature);
75        self
76    }
77
78    pub fn with_stop_words(mut self, stop_words: Vec<String>) -> Self {
79        self.stop_words = Some(stop_words);
80        self
81    }
82
83    //TODO:Check if this should be a &str instead of a String
84    pub fn with_streaming_func<F, Fut>(mut self, mut func: F) -> Self
85    where
86        F: FnMut(String) -> Fut + Send + 'static,
87        Fut: Future<Output = Result<(), ()>> + Send + 'static,
88    {
89        let func = Arc::new(Mutex::new(
90            move |s: String| -> Pin<Box<dyn Future<Output = Result<(), ()>> + Send>> {
91                Box::pin(func(s))
92            },
93        ));
94
95        self.streaming_func = Some(func);
96        self
97    }
98
99    pub fn with_top_k(mut self, top_k: usize) -> Self {
100        self.top_k = Some(top_k);
101        self
102    }
103
104    pub fn with_top_p(mut self, top_p: f32) -> Self {
105        self.top_p = Some(top_p);
106        self
107    }
108
109    pub fn with_seed(mut self, seed: usize) -> Self {
110        self.seed = Some(seed);
111        self
112    }
113
114    pub fn with_min_length(mut self, min_length: usize) -> Self {
115        self.min_length = Some(min_length);
116        self
117    }
118
119    pub fn with_max_length(mut self, max_length: usize) -> Self {
120        self.max_length = Some(max_length);
121        self
122    }
123
124    pub fn with_n(mut self, n: usize) -> Self {
125        self.n = Some(n);
126        self
127    }
128
129    pub fn with_repetition_penalty(mut self, repetition_penalty: f32) -> Self {
130        self.repetition_penalty = Some(repetition_penalty);
131        self
132    }
133
134    pub fn with_frequency_penalty(mut self, frequency_penalty: f32) -> Self {
135        self.frequency_penalty = Some(frequency_penalty);
136        self
137    }
138
139    pub fn with_presence_penalty(mut self, presence_penalty: f32) -> Self {
140        self.presence_penalty = Some(presence_penalty);
141        self
142    }
143
144    pub fn with_functions(mut self, functions: Vec<FunctionDefinition>) -> Self {
145        self.functions = Some(functions);
146        self
147    }
148
149    pub fn with_function_call_behavior(mut self, behavior: FunctionCallBehavior) -> Self {
150        self.function_call_behavior = Some(behavior);
151        self
152    }
153
154    pub fn with_response_format(mut self, response_format: ResponseFormat) -> Self {
155        self.response_format = Some(response_format);
156        self
157    }
158
159    pub fn with_stream_usage(mut self, stream_usage: bool) -> Self {
160        self.stream_usage = Some(stream_usage);
161        self
162    }
163
164    pub fn merge_options(&mut self, incoming_options: CallOptions) {
165        // For simple scalar types wrapped in Option, prefer incoming option if it is Some
166        self.candidate_count = incoming_options.candidate_count.or(self.candidate_count);
167        self.max_tokens = incoming_options.max_tokens.or(self.max_tokens);
168        self.temperature = incoming_options.temperature.or(self.temperature);
169        self.top_k = incoming_options.top_k.or(self.top_k);
170        self.top_p = incoming_options.top_p.or(self.top_p);
171        self.seed = incoming_options.seed.or(self.seed);
172        self.min_length = incoming_options.min_length.or(self.min_length);
173        self.max_length = incoming_options.max_length.or(self.max_length);
174        self.n = incoming_options.n.or(self.n);
175        self.repetition_penalty = incoming_options
176            .repetition_penalty
177            .or(self.repetition_penalty);
178        self.frequency_penalty = incoming_options
179            .frequency_penalty
180            .or(self.frequency_penalty);
181        self.presence_penalty = incoming_options.presence_penalty.or(self.presence_penalty);
182        self.function_call_behavior = incoming_options
183            .function_call_behavior
184            .or(self.function_call_behavior.clone());
185        self.response_format = incoming_options
186            .response_format
187            .or(self.response_format.clone());
188        self.stream_usage = incoming_options.stream_usage.or(self.stream_usage);
189
190        // For `Vec<String>`, merge if both are Some; prefer incoming if only incoming is Some
191        if let Some(mut new_stop_words) = incoming_options.stop_words {
192            if let Some(existing_stop_words) = &mut self.stop_words {
193                existing_stop_words.append(&mut new_stop_words);
194            } else {
195                self.stop_words = Some(new_stop_words);
196            }
197        }
198
199        // For `Vec<FunctionDefinition>`, similar logic to `Vec<String>`
200        if let Some(mut incoming_functions) = incoming_options.functions {
201            if let Some(existing_functions) = &mut self.functions {
202                existing_functions.append(&mut incoming_functions);
203            } else {
204                self.functions = Some(incoming_functions);
205            }
206        }
207
208        // `streaming_func` requires a judgment call on how you want to handle merging.
209        // Here, the incoming option simply replaces the existing one if it's Some.
210        self.streaming_func = incoming_options
211            .streaming_func
212            .or_else(|| self.streaming_func.clone());
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use serde_json::json;
220
221    #[test]
222    fn merge_prefers_incoming_scalars_and_keeps_existing() {
223        let mut base = CallOptions::new()
224            .with_max_tokens(100)
225            .with_temperature(0.1)
226            .with_seed(7);
227        base.merge_options(CallOptions::new().with_temperature(0.9).with_top_p(0.5));
228
229        assert_eq!(base.max_tokens, Some(100), "kept when incoming is None");
230        assert_eq!(base.temperature, Some(0.9), "incoming wins");
231        assert_eq!(base.top_p, Some(0.5), "new field set");
232        assert_eq!(base.seed, Some(7));
233    }
234
235    #[test]
236    fn merge_appends_stop_words_and_functions() {
237        let f = |n: &str| FunctionDefinition::new(n, "d", json!({}));
238        let mut base = CallOptions::new()
239            .with_stop_words(vec!["a".into()])
240            .with_functions(vec![f("one")]);
241        base.merge_options(
242            CallOptions::new()
243                .with_stop_words(vec!["b".into()])
244                .with_functions(vec![f("two")]),
245        );
246        assert_eq!(base.stop_words, Some(vec!["a".into(), "b".into()]));
247        let names: Vec<_> = base
248            .functions
249            .unwrap()
250            .into_iter()
251            .map(|f| f.name)
252            .collect();
253        assert_eq!(names, vec!["one", "two"]);
254
255        let mut empty = CallOptions::new();
256        empty.merge_options(CallOptions::new().with_stop_words(vec!["x".into()]));
257        assert_eq!(empty.stop_words, Some(vec!["x".into()]));
258    }
259
260    #[test]
261    fn merge_streaming_func_and_behaviors() {
262        let mut base = CallOptions::new().with_streaming_func(|_s| async { Ok(()) });
263        base.merge_options(
264            CallOptions::new()
265                .with_function_call_behavior(FunctionCallBehavior::Auto)
266                .with_response_format(ResponseFormat::JsonObject)
267                .with_stream_usage(true),
268        );
269        assert!(base.streaming_func.is_some(), "existing func kept");
270        assert!(matches!(
271            base.function_call_behavior,
272            Some(FunctionCallBehavior::Auto)
273        ));
274        assert!(matches!(
275            base.response_format,
276            Some(ResponseFormat::JsonObject)
277        ));
278        assert_eq!(base.stream_usage, Some(true));
279    }
280
281    #[tokio::test]
282    async fn streaming_func_is_invoked() {
283        let seen = Arc::new(Mutex::new(Vec::<String>::new()));
284        let seen_clone = seen.clone();
285        let opts = CallOptions::new().with_streaming_func(move |s| {
286            let seen = seen_clone.clone();
287            async move {
288                seen.lock().await.push(s);
289                Ok(())
290            }
291        });
292        let func = opts.streaming_func.unwrap();
293        let mut f = func.lock().await;
294        f("chunk".to_string()).await.unwrap();
295        assert_eq!(*seen.lock().await, vec!["chunk".to_string()]);
296    }
297}
langchain_rust/language_models/options.rs

langchain_rust/language_models/
options.rs
