allms/llm_models/
tools.rs

1use log::warn;
2use serde::{Deserialize, Serialize};
3use serde_json::{json, to_value, Value};
4
5///
6/// Enum of all the tools that can be used with different LLM providers
7///
8#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
9pub enum LLMTools {
10    /// OpenAI
11    OpenAIFileSearch(OpenAIFileSearchConfig),
12    OpenAIWebSearch(OpenAIWebSearchConfig),
13    OpenAIComputerUse(OpenAIComputerUseConfig),
14    OpenAIReasoning(OpenAIReasoningConfig),
15    OpenAICodeInterpreter(OpenAICodeInterpreterConfig),
16    /// Anthropic
17    AnthropicCodeExecution(AnthropicCodeExecutionConfig),
18    AnthropicComputerUse(AnthropicComputerUseConfig),
19    AnthropicFileSearch(AnthropicFileSearchConfig),
20    AnthropicWebSearch(AnthropicWebSearchConfig),
21    /// xAI
22    XAIWebSearch(XAIWebSearchConfig),
23    XAIXSearch(XAIXSearchConfig),
24    /// Gemini
25    GeminiCodeInterpreter(GeminiCodeInterpreterConfig),
26    GeminiWebSearch(GeminiWebSearchConfig),
27    /// Mistral
28    MistralWebSearch(MistralWebSearchConfig),
29    MistralCodeInterpreter(MistralCodeInterpreterConfig),
30}
31
32impl LLMTools {
33    pub fn get_config_json(&self) -> Option<Value> {
34        match self {
35            LLMTools::OpenAIFileSearch(cfg) => to_value(cfg).ok(),
36            LLMTools::OpenAIWebSearch(cfg) => to_value(cfg).ok(),
37            LLMTools::OpenAIComputerUse(cfg) => to_value(cfg).ok(),
38            LLMTools::OpenAIReasoning(cfg) => to_value(cfg).ok(),
39            LLMTools::OpenAICodeInterpreter(cfg) => to_value(cfg).ok(),
40            LLMTools::AnthropicCodeExecution(cfg) => to_value(cfg).ok(),
41            LLMTools::AnthropicComputerUse(cfg) => to_value(cfg).ok(),
42            LLMTools::AnthropicFileSearch(cfg) => to_value(cfg).ok(),
43            LLMTools::AnthropicWebSearch(cfg) => to_value(cfg).ok(),
44            LLMTools::XAIWebSearch(cfg) => to_value(cfg).ok(),
45            LLMTools::XAIXSearch(cfg) => to_value(cfg).ok(),
46            LLMTools::GeminiCodeInterpreter(cfg) => to_value(cfg).ok(),
47            // For Gemini Web Search we decode configuration based on the settings
48            LLMTools::GeminiWebSearch(cfg) => Some(cfg.get_config_json()),
49            LLMTools::MistralWebSearch(cfg) => to_value(cfg).ok(),
50            LLMTools::MistralCodeInterpreter(cfg) => to_value(cfg).ok(),
51        }
52    }
53}
54
55///
56/// OpenAI File Search tool config
57///
58#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
59pub struct OpenAIFileSearchConfig {
60    #[serde(rename = "type")]
61    pub tool_type: OpenAIFileSearchToolType,
62    pub vector_store_ids: Vec<String>,
63    pub max_num_results: Option<usize>,
64}
65
66impl OpenAIFileSearchConfig {
67    pub fn new(vector_store_ids: Vec<String>) -> Self {
68        Self {
69            tool_type: OpenAIFileSearchToolType::FileSearch,
70            vector_store_ids,
71            max_num_results: None,
72        }
73    }
74}
75
76#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
77pub enum OpenAIFileSearchToolType {
78    #[serde(rename = "file_search")]
79    FileSearch,
80}
81
82///
83/// OpenAI Web Search tool config
84///
85#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
86pub struct OpenAIWebSearchConfig {
87    #[serde(rename = "type")]
88    pub tool_type: OpenAIWebSearchToolType,
89    pub search_context_size: Option<OpenAIWebSearchContextSize>,
90}
91
92impl Default for OpenAIWebSearchConfig {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98impl OpenAIWebSearchConfig {
99    pub fn new() -> Self {
100        Self {
101            tool_type: OpenAIWebSearchToolType::default(),
102            search_context_size: None,
103        }
104    }
105}
106
107#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
108pub enum OpenAIWebSearchToolType {
109    #[serde(rename = "web_search_preview")]
110    #[default]
111    WebSearchPreview,
112    #[serde(rename = "web_search_preview_2025_03_11")]
113    WebSearchPreview20250311,
114}
115
116#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
117pub enum OpenAIWebSearchContextSize {
118    #[serde(rename = "low")]
119    Low,
120    #[serde(rename = "medium")]
121    Medium,
122    #[serde(rename = "high")]
123    High,
124}
125
126///
127/// OpenAI Computer Use tool config
128///
129#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
130pub struct OpenAIComputerUseConfig {
131    #[serde(rename = "type")]
132    pub tool_type: OpenAIComputerUseToolType,
133    pub display_height: usize,
134    pub display_width: usize,
135    pub environment: String,
136}
137
138#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
139pub enum OpenAIComputerUseToolType {
140    #[serde(rename = "computer_use_preview")]
141    ComputerUsePreview,
142}
143
144impl OpenAIComputerUseConfig {
145    pub fn new(display_height: usize, display_width: usize, environment: String) -> Self {
146        Self {
147            tool_type: OpenAIComputerUseToolType::ComputerUsePreview,
148            display_height,
149            display_width,
150            environment,
151        }
152    }
153}
154
155///
156/// OpenAI Reasoning config
157///
158#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
159pub struct OpenAIReasoningConfig {
160    pub effort: Option<OpenAIReasoningEffort>,
161    pub summary: Option<OpenAIReasoningSummary>,
162}
163
164impl OpenAIReasoningConfig {
165    pub fn new(
166        effort: Option<OpenAIReasoningEffort>,
167        summary: Option<OpenAIReasoningSummary>,
168    ) -> Self {
169        Self { effort, summary }
170    }
171}
172
173#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
174pub enum OpenAIReasoningEffort {
175    #[serde(rename = "low")]
176    Low,
177    #[serde(rename = "medium")]
178    Medium,
179    #[serde(rename = "high")]
180    High,
181}
182
183#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
184pub enum OpenAIReasoningSummary {
185    #[serde(rename = "auto")]
186    Auto,
187    #[serde(rename = "none")]
188    Concise,
189    #[serde(rename = "detailed")]
190    Detailed,
191}
192
193///
194/// OpenAI Code Interpreter tool config
195///
196#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
197pub struct OpenAICodeInterpreterConfig {
198    #[serde(rename = "type")]
199    pub tool_type: OpenAICodeInterpreterToolType,
200    pub container: OpenAICodeInterpreterContainerConfig,
201}
202
203#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
204pub enum OpenAICodeInterpreterToolType {
205    #[serde(rename = "code_interpreter")]
206    CodeInterpreter,
207}
208
209// Can be a container ID or an object that specifies uploaded file IDs to make available to your code.
210#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
211#[serde(untagged)]
212pub enum OpenAICodeInterpreterContainerConfig {
213    ContainerId(String),
214    CodeInterpreterContainerAuto(CodeInterpreterContainerAutoConfig),
215}
216
217#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
218pub struct CodeInterpreterContainerAutoConfig {
219    #[serde(rename = "type")]
220    container_type: OpenAICodeInterpreterContainerType,
221    file_ids: Vec<String>,
222}
223
224#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
225pub enum OpenAICodeInterpreterContainerType {
226    #[serde(rename = "auto")]
227    #[default]
228    Auto,
229}
230
231impl Default for OpenAICodeInterpreterConfig {
232    fn default() -> Self {
233        Self::new()
234    }
235}
236
237impl OpenAICodeInterpreterConfig {
238    pub fn new() -> Self {
239        Self {
240            tool_type: OpenAICodeInterpreterToolType::CodeInterpreter,
241            container: OpenAICodeInterpreterContainerConfig::CodeInterpreterContainerAuto(
242                CodeInterpreterContainerAutoConfig::default(),
243            ),
244        }
245    }
246}
247
248///
249/// Anthropic Code Execution tool config
250///
251#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
252pub struct AnthropicCodeExecutionConfig {
253    pub name: AnthropicCodeExecutionName,
254    #[serde(rename = "type")]
255    pub tool_type: AnthropicCodeExecutionToolType,
256    pub cache_control: Option<AnthropicCacheControl>,
257}
258
259#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
260pub enum AnthropicCodeExecutionName {
261    #[serde(rename = "code_execution")]
262    #[default]
263    CodeExecution,
264}
265
266#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
267pub enum AnthropicCodeExecutionToolType {
268    #[serde(rename = "code_execution_20250522")]
269    #[default]
270    CodeExecution20250522,
271    #[serde(rename = "code_execution_20260120")]
272    CodeExecution20260120,
273}
274
275#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
276pub struct AnthropicCacheControl {
277    #[serde(rename = "type")]
278    pub cache_type: AnthropicCacheControlType,
279    pub ttl: AnthropicCacheControllTTL,
280}
281
282#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
283pub enum AnthropicCacheControlType {
284    #[serde(rename = "ephemeral")]
285    #[default]
286    Ephemeral,
287}
288
289#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
290pub enum AnthropicCacheControllTTL {
291    #[serde(rename = "5m")]
292    #[default]
293    FiveMinutes,
294    #[serde(rename = "1h")]
295    OneHour,
296}
297
298impl AnthropicCodeExecutionConfig {
299    pub fn new() -> Self {
300        Self {
301            name: AnthropicCodeExecutionName::default(),
302            tool_type: AnthropicCodeExecutionToolType::default(),
303            cache_control: None,
304        }
305    }
306
307    pub fn set_type(mut self, tool_type: AnthropicCodeExecutionToolType) -> Self {
308        self.tool_type = tool_type;
309        self
310    }
311
312    pub fn cache_control(mut self, ttl: AnthropicCacheControllTTL) -> Self {
313        self.cache_control = Some(AnthropicCacheControl {
314            cache_type: AnthropicCacheControlType::default(),
315            ttl,
316        });
317        self
318    }
319}
320
321///
322/// Anthropic Computer Use tool config
323///
324#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
325pub struct AnthropicComputerUseConfig {
326    pub name: AnthropicComputerUseName,
327    #[serde(rename = "type")]
328    pub tool_type: AnthropicComputerUseToolType,
329    pub display_height_px: usize,
330    pub display_width_px: usize,
331    pub cache_control: Option<AnthropicCacheControl>,
332    pub display_number: Option<usize>,
333}
334
335#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
336pub enum AnthropicComputerUseName {
337    #[serde(rename = "computer")]
338    #[default]
339    Computer,
340}
341
342#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
343pub enum AnthropicComputerUseToolType {
344    #[serde(rename = "computer_20241022")]
345    #[default]
346    Computer20241022,
347}
348
349impl AnthropicComputerUseConfig {
350    pub fn new(display_height_px: usize, display_width_px: usize) -> Self {
351        Self {
352            name: AnthropicComputerUseName::default(),
353            tool_type: AnthropicComputerUseToolType::default(),
354            display_height_px: display_height_px.max(1),
355            display_width_px: display_width_px.max(1),
356            cache_control: None,
357            display_number: None,
358        }
359    }
360
361    pub fn cache_control(mut self, ttl: AnthropicCacheControllTTL) -> Self {
362        self.cache_control = Some(AnthropicCacheControl {
363            cache_type: AnthropicCacheControlType::default(),
364            ttl,
365        });
366        self
367    }
368
369    pub fn display_number(mut self, display_number: usize) -> Self {
370        self.display_number = Some(display_number);
371        self
372    }
373}
374
375///
376/// Anthropic File Search tool config
377///
378#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
379pub struct AnthropicFileSearchConfig {
380    pub file_id: String,
381}
382
383impl AnthropicFileSearchConfig {
384    pub fn new(file_id: String) -> Self {
385        Self { file_id }
386    }
387
388    pub fn content(&self) -> Value {
389        json!({
390            "type": "document",
391            "source": {
392                "type": "file",
393                "file_id": self.file_id,
394            },
395        })
396    }
397}
398
399///
400/// xAI Web Search tool config
401///
402/// This config directly matches the xAI API format for web search tools.
403#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
404pub struct XAIWebSearchConfig {
405    #[serde(rename = "type")]
406    pub tool_type: XAIWebSearchToolType,
407    pub filters: Option<XAIWebSearchFilters>,
408    pub enable_image_understanding: Option<bool>,
409}
410
411#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
412pub enum XAIWebSearchToolType {
413    #[serde(rename = "web_search")]
414    #[default]
415    WebSearch,
416}
417
418#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
419pub struct XAIWebSearchFilters {
420    pub allowed_domains: Option<Vec<String>>,
421    pub excluded_domains: Option<Vec<String>>,
422}
423
424impl Default for XAIWebSearchConfig {
425    fn default() -> Self {
426        Self::new()
427    }
428}
429
430impl XAIWebSearchConfig {
431    pub fn new() -> Self {
432        Self {
433            tool_type: XAIWebSearchToolType::WebSearch,
434            filters: None,
435            enable_image_understanding: None,
436        }
437    }
438
439    /// Add allowed domains to the web search filters
440    pub fn add_allowed_domains(mut self, domains: &[String]) -> Self {
441        if domains.is_empty() {
442            return self;
443        }
444        // xAI API only supports allowed domains or excluded domains, so we will clear any excluded domains
445        if self
446            .filters
447            .as_ref()
448            .and_then(|f| f.excluded_domains.as_ref())
449            .is_some()
450        {
451            warn!("[allms][xAI][Tools] Adding allowed domains will clear any excluded domains");
452        }
453        let current_filters = self.filters.unwrap_or_default();
454        let mut allowed_domains = current_filters.allowed_domains.unwrap_or_default();
455        allowed_domains.extend(domains.to_vec());
456        self.filters = Some(XAIWebSearchFilters {
457            allowed_domains: Some(allowed_domains),
458            excluded_domains: None, // Clear excluded domains when adding allowed domains
459        });
460        self
461    }
462
463    /// Add excluded domains to the web search filters
464    pub fn add_excluded_domains(mut self, domains: &[String]) -> Self {
465        if domains.is_empty() {
466            return self;
467        }
468        // xAI API only supports allowed domains or excluded domains, so we will clear any allowed domains
469        if self
470            .filters
471            .as_ref()
472            .and_then(|f| f.allowed_domains.as_ref())
473            .is_some()
474        {
475            warn!("[allms][xAI][Tools] Adding excluded domains will clear any allowed domains");
476        }
477        let current_filters = self.filters.unwrap_or_default();
478        let mut excluded_domains = current_filters.excluded_domains.unwrap_or_default();
479        excluded_domains.extend(domains.to_vec());
480        self.filters = Some(XAIWebSearchFilters {
481            allowed_domains: None, // Clear allowed domains when adding excluded domains
482            excluded_domains: Some(excluded_domains),
483        });
484        self
485    }
486
487    pub fn with_enable_image_understanding(mut self, enable: bool) -> Self {
488        self.enable_image_understanding = Some(enable);
489        self
490    }
491}
492
493///
494/// xAI X Search tool config
495///
496/// This config directly matches the xAI API format for X search tools.
497#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
498pub struct XAIXSearchConfig {
499    #[serde(rename = "type")]
500    pub tool_type: XAIXSearchToolType,
501    pub allowed_x_handles: Option<Vec<String>>,
502    pub excluded_x_handles: Option<Vec<String>>,
503    pub from_date: Option<String>,
504    pub to_date: Option<String>,
505    pub enable_image_understanding: Option<bool>,
506    pub enable_video_understanding: Option<bool>,
507}
508
509#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
510pub enum XAIXSearchToolType {
511    #[serde(rename = "x_search")]
512    #[default]
513    XSearch,
514}
515
516impl Default for XAIXSearchConfig {
517    fn default() -> Self {
518        Self::new()
519    }
520}
521
522impl XAIXSearchConfig {
523    pub fn new() -> Self {
524        Self {
525            tool_type: XAIXSearchToolType::XSearch,
526            allowed_x_handles: None,
527            excluded_x_handles: None,
528            from_date: None,
529            to_date: None,
530            enable_image_understanding: None,
531            enable_video_understanding: None,
532        }
533    }
534
535    /// Add allowed X handles to the X search filters
536    pub fn add_allowed_x_handles(mut self, handles: &[String]) -> Self {
537        if handles.is_empty() {
538            return self;
539        }
540        // xAI API only supports allowed handles or excluded handles, so we will clear any excluded handles
541        if self.excluded_x_handles.is_some() {
542            warn!("[allms][xAI][Tools] Adding allowed X handles will clear any excluded X handles");
543        }
544        let mut allowed_handles = self.allowed_x_handles.unwrap_or_default();
545        allowed_handles.extend(handles.to_vec());
546        self.allowed_x_handles = Some(allowed_handles);
547        self.excluded_x_handles = None; // Clear excluded handles when adding allowed handles
548        self
549    }
550
551    /// Add excluded X handles to the X search filters
552    pub fn add_excluded_x_handles(mut self, handles: &[String]) -> Self {
553        if handles.is_empty() {
554            return self;
555        }
556        // xAI API only supports allowed handles or excluded handles, so we will clear any allowed handles
557        if self.allowed_x_handles.is_some() {
558            warn!("[allms][xAI][Tools] Adding excluded X handles will clear any allowed X handles");
559        }
560        let mut excluded_handles = self.excluded_x_handles.unwrap_or_default();
561        excluded_handles.extend(handles.to_vec());
562        self.excluded_x_handles = Some(excluded_handles);
563        self.allowed_x_handles = None; // Clear allowed handles when adding excluded handles
564        self
565    }
566
567    pub fn from_date(mut self, from_date: String) -> Self {
568        self.from_date = Some(from_date);
569        self
570    }
571
572    pub fn to_date(mut self, to_date: String) -> Self {
573        self.to_date = Some(to_date);
574        self
575    }
576
577    pub fn enable_image_understanding(mut self, enable: bool) -> Self {
578        self.enable_image_understanding = Some(enable);
579        self
580    }
581
582    pub fn enable_video_understanding(mut self, enable: bool) -> Self {
583        self.enable_video_understanding = Some(enable);
584        self
585    }
586}
587
588///
589/// Anthropic Web Search tool config
590///
591#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
592pub struct AnthropicWebSearchConfig {
593    pub name: AnthropicWebSearchName,
594    #[serde(rename = "type")]
595    pub tool_type: AnthropicWebSearchToolType,
596    pub allowed_domains: Option<Vec<String>>,
597    pub blocked_domains: Option<Vec<String>>,
598    pub cache_control: Option<AnthropicCacheControl>,
599    pub max_uses: Option<usize>,
600    pub user_location: Option<AnthropicWebSearchUserLocation>,
601}
602
603#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
604pub enum AnthropicWebSearchName {
605    #[serde(rename = "web_search")]
606    #[default]
607    WebSearch,
608}
609
610#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
611pub enum AnthropicWebSearchToolType {
612    #[serde(rename = "web_search_20250305")]
613    #[default]
614    WebSearch20250305,
615    #[serde(rename = "web_search_20260209")]
616    WebSearch20260209,
617}
618
619#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
620pub struct AnthropicWebSearchUserLocation {
621    #[serde(rename = "type")]
622    pub location_type: AnthropicWebSearchUserLocationType,
623    pub city: Option<String>,
624    pub country: Option<String>,
625    pub region: Option<String>,
626    pub timezone: Option<String>,
627}
628
629#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
630pub enum AnthropicWebSearchUserLocationType {
631    #[serde(rename = "approximate")]
632    #[default]
633    Approximate,
634}
635
636impl AnthropicWebSearchConfig {
637    pub fn new() -> Self {
638        Self {
639            name: AnthropicWebSearchName::default(),
640            tool_type: AnthropicWebSearchToolType::default(),
641            allowed_domains: None,
642            blocked_domains: None,
643            cache_control: None,
644            max_uses: None,
645            user_location: None,
646        }
647    }
648
649    pub fn set_type(mut self, tool_type: AnthropicWebSearchToolType) -> Self {
650        self.tool_type = tool_type;
651        self
652    }
653
654    pub fn cache_control(mut self, ttl: AnthropicCacheControllTTL) -> Self {
655        self.cache_control = Some(AnthropicCacheControl {
656            cache_type: AnthropicCacheControlType::default(),
657            ttl,
658        });
659        self
660    }
661
662    pub fn allowed_domains(mut self, allowed_domains: Vec<String>) -> Self {
663        self.allowed_domains = Some(allowed_domains);
664        if self.blocked_domains.is_some() {
665            warn!("[allms][Anthropic][Tools] Allowed domains will clear any blocked domains");
666            self.blocked_domains = None;
667        }
668        self
669    }
670
671    pub fn blocked_domains(mut self, blocked_domains: Vec<String>) -> Self {
672        self.blocked_domains = Some(blocked_domains);
673        if self.allowed_domains.is_some() {
674            warn!("[allms][Anthropic][Tools] Blocked domains will clear any allowed domains");
675            self.allowed_domains = None;
676        }
677        self
678    }
679
680    pub fn max_uses(mut self, max_uses: usize) -> Self {
681        self.max_uses = Some(max_uses);
682        self
683    }
684
685    pub fn user_location(mut self, user_location: AnthropicWebSearchUserLocation) -> Self {
686        self.user_location = Some(user_location);
687        self
688    }
689}
690
691impl AnthropicWebSearchUserLocation {
692    pub fn new(location_type: AnthropicWebSearchUserLocationType) -> Self {
693        Self {
694            location_type,
695            city: None,
696            country: None,
697            region: None,
698            timezone: None,
699        }
700    }
701
702    pub fn city(mut self, city: String) -> Self {
703        self.city = Some(city);
704        self
705    }
706
707    pub fn country(mut self, country: String) -> Self {
708        self.country = Some(country);
709        self
710    }
711
712    pub fn region(mut self, region: String) -> Self {
713        self.region = Some(region);
714        self
715    }
716
717    pub fn timezone(mut self, timezone: String) -> Self {
718        self.timezone = Some(timezone);
719        self
720    }
721}
722
723///
724/// Gemini Code Interpreter
725///
726#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
727pub struct GeminiCodeInterpreterConfig {
728    pub code_execution: GeminiCodeExecutionTool,
729}
730
731#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
732pub struct GeminiCodeExecutionTool {}
733
734impl GeminiCodeInterpreterConfig {
735    pub fn new() -> Self {
736        Self {
737            code_execution: GeminiCodeExecutionTool {},
738        }
739    }
740}
741
742///
743/// Gemini Web Search
744///
745#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq)]
746pub struct GeminiWebSearchConfig {
747    context_urls: Vec<String>,
748    include_web: bool,
749}
750
751impl Default for GeminiWebSearchConfig {
752    fn default() -> Self {
753        Self::new()
754    }
755}
756
757impl GeminiWebSearchConfig {
758    pub fn new() -> Self {
759        Self {
760            context_urls: Vec::new(),
761            include_web: false,
762        }
763    }
764
765    /// Add a single URL to the context URLs list
766    pub fn add_source(mut self, url: &str) -> Self {
767        self.context_urls.push(url.to_string());
768        self
769    }
770
771    /// Add multiple URLs to the context URLs list
772    pub fn add_sources(mut self, urls: &[String]) -> Self {
773        self.context_urls.extend(urls.to_vec());
774        self
775    }
776
777    /// Enable google search in addition to URL context
778    pub fn include_web(mut self) -> Self {
779        self.include_web = true;
780        self
781    }
782
783    /// Get the list of context URLs
784    pub fn get_context_urls(&self) -> &[String] {
785        &self.context_urls
786    }
787
788    /// Get the configuration as JSON based on the current state
789    pub fn get_config_json(&self) -> serde_json::Value {
790        if self.context_urls.is_empty() {
791            // Case A: context_urls is empty
792            serde_json::json!({
793                "google_search": {}
794            })
795        } else if !self.include_web {
796            // Case B: context_urls is not empty and include_web is false
797            serde_json::json!({
798                "url_context": {}
799            })
800        } else {
801            // Case C: context_urls is not empty and include_web is true
802            serde_json::json!([
803                {
804                    "url_context": {}
805                },
806                {
807                    "google_search": {}
808                }
809            ])
810        }
811    }
812}
813
814///
815/// Mistral Web Search
816///
817#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
818pub struct MistralWebSearchConfig {
819    #[serde(rename = "type")]
820    pub web_search_type: MistralWebSearchType,
821}
822
823#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
824pub enum MistralWebSearchType {
825    #[serde(rename = "web_search")]
826    #[default]
827    WebSearch,
828    #[serde(rename = "web_search_premium")]
829    WebSearchPremium,
830}
831
832impl MistralWebSearchConfig {
833    pub fn new() -> Self {
834        Self {
835            web_search_type: MistralWebSearchType::default(),
836        }
837    }
838
839    pub fn set_type(mut self, web_search_type: MistralWebSearchType) -> Self {
840        self.web_search_type = web_search_type;
841        self
842    }
843
844    pub fn get_type_str(&self) -> String {
845        serde_json::to_string(&self.web_search_type).unwrap_or_default()
846    }
847}
848
849///
850/// Mistral Code Interpreter
851///
852#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
853pub struct MistralCodeInterpreterConfig {
854    #[serde(rename = "type")]
855    pub code_interpreter_type: MistralCodeInterpreterType,
856}
857
858#[derive(Deserialize, Serialize, Debug, Clone, Eq, PartialEq, Default)]
859pub enum MistralCodeInterpreterType {
860    #[serde(rename = "code_interpreter")]
861    #[default]
862    CodeInterpreter,
863}
864
865impl MistralCodeInterpreterConfig {
866    pub fn new() -> Self {
867        Self {
868            code_interpreter_type: MistralCodeInterpreterType::default(),
869        }
870    }
871}
872
873///
874/// Tests
875///
876#[cfg(test)]
877mod tests {
878    use super::*;
879
880    #[test]
881    fn test_gemini_web_search_config_empty() {
882        let config = GeminiWebSearchConfig::new();
883        let json = config.get_config_json();
884
885        // Case A: context_urls is empty -> should return google_search only
886        assert_eq!(json["google_search"], serde_json::json!({}));
887        assert!(
888            json["url_context"].is_null(),
889            "Should not have url_context when empty"
890        );
891        assert!(!json.is_array(), "Should not be an array when empty");
892    }
893
894    #[test]
895    fn test_gemini_web_search_config_with_urls() {
896        let config = GeminiWebSearchConfig::new().add_source("https://example.com");
897        let json = config.get_config_json();
898
899        // Case B: context_urls is not empty and include_web is false -> should return url_context only
900        assert_eq!(json["url_context"], serde_json::json!({}));
901        assert!(
902            json["google_search"].is_null(),
903            "Should not have google_search when only URLs"
904        );
905        assert!(!json.is_array(), "Should not be an array when only URLs");
906    }
907
908    #[test]
909    fn test_gemini_web_search_config_with_urls_and_web() {
910        let config = GeminiWebSearchConfig::new()
911            .add_source("https://example.com")
912            .include_web();
913        let json = config.get_config_json();
914
915        // Case C: context_urls is not empty and include_web is true -> should return array with both
916        assert!(json.is_array());
917        assert_eq!(json[0]["url_context"], serde_json::json!({}));
918        assert_eq!(json[1]["google_search"], serde_json::json!({}));
919    }
920
921    #[test]
922    fn test_gemini_web_search_config_multiple_sources() {
923        let config = GeminiWebSearchConfig::new().add_sources(&[
924            "https://site1.com".to_string(),
925            "https://site2.com".to_string(),
926        ]);
927        let json = config.get_config_json();
928
929        // Should return url_context for multiple sources without include_web
930        assert_eq!(json["url_context"], serde_json::json!({}));
931        assert!(
932            json["google_search"].is_null(),
933            "Should not have google_search when only URLs"
934        );
935        assert!(!json.is_array(), "Should not be an array when only URLs");
936    }
937
938    #[test]
939    fn test_gemini_web_search_config_builder_pattern() {
940        let config = GeminiWebSearchConfig::new()
941            .add_source("https://example.com")
942            .add_sources(&["https://site1.com".to_string()])
943            .include_web();
944
945        // Verify the internal state
946        assert_eq!(config.get_context_urls().len(), 2);
947        assert!(config.include_web);
948
949        let json = config.get_config_json();
950        assert!(json.is_array());
951        assert_eq!(json[0]["url_context"], serde_json::json!({}));
952        assert_eq!(json[1]["google_search"], serde_json::json!({}));
953    }
954
955    #[test]
956    fn test_gemini_web_search_config_through_llm_tools() {
957        let tool = LLMTools::GeminiWebSearch(
958            GeminiWebSearchConfig::new().add_source("https://example.com"),
959        );
960
961        let json = tool.get_config_json().unwrap();
962        assert_eq!(json["url_context"], serde_json::json!({}));
963    }
964}
allms/llm_models/tools.rs

allms/llm_models/
tools.rs
