openai_ergonomic/builders/
assistants.rs

1//! Assistants API builders.
2//!
3//! This module provides ergonomic builders for `OpenAI` Assistants API operations,
4//! including creating assistants, managing threads, messages, and runs.
5//!
6//! Note: This is a simplified implementation focusing on the most commonly used features.
7
8use crate::Result;
9use openai_client_base::models;
10use serde_json::Value;
11use std::collections::HashMap;
12
13/// Builder for creating a new assistant.
14///
15/// This builder provides a fluent interface for creating `OpenAI` assistants
16/// with commonly used parameters including tool support.
17#[derive(Debug, Clone)]
18pub struct AssistantBuilder {
19    model: String,
20    name: Option<String>,
21    description: Option<String>,
22    instructions: Option<String>,
23    tools: Vec<AssistantTool>,
24    metadata: HashMap<String, String>,
25}
26
27impl AssistantBuilder {
28    /// Create a new assistant builder with the specified model.
29    ///
30    /// # Examples
31    ///
32    /// ```rust
33    /// use openai_ergonomic::builders::assistants::{AssistantBuilder, tool_code_interpreter};
34    ///
35    /// let builder = AssistantBuilder::new("gpt-4")
36    ///     .name("My Assistant")
37    ///     .instructions("You are a helpful coding assistant.")
38    ///     .add_tool(tool_code_interpreter());
39    /// ```
40    #[must_use]
41    pub fn new(model: impl Into<String>) -> Self {
42        Self {
43            model: model.into(),
44            name: None,
45            description: None,
46            instructions: None,
47            tools: Vec::new(),
48            metadata: HashMap::new(),
49        }
50    }
51
52    /// Set the assistant's name.
53    #[must_use]
54    pub fn name(mut self, name: impl Into<String>) -> Self {
55        self.name = Some(name.into());
56        self
57    }
58
59    /// Set the assistant's description.
60    #[must_use]
61    pub fn description(mut self, description: impl Into<String>) -> Self {
62        self.description = Some(description.into());
63        self
64    }
65
66    /// Set the assistant's instructions (system prompt).
67    #[must_use]
68    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
69        self.instructions = Some(instructions.into());
70        self
71    }
72
73    /// Add tools to the assistant.
74    #[must_use]
75    pub fn tools(mut self, tools: Vec<AssistantTool>) -> Self {
76        self.tools = tools;
77        self
78    }
79
80    /// Add a single tool to the assistant.
81    #[must_use]
82    pub fn add_tool(mut self, tool: AssistantTool) -> Self {
83        self.tools.push(tool);
84        self
85    }
86
87    /// Add metadata to the assistant.
88    #[must_use]
89    pub fn metadata(mut self, metadata: HashMap<String, String>) -> Self {
90        self.metadata = metadata;
91        self
92    }
93
94    /// Add a single metadata key-value pair.
95    #[must_use]
96    pub fn add_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
97        self.metadata.insert(key.into(), value.into());
98        self
99    }
100
101    /// Get the model for this assistant.
102    #[must_use]
103    pub fn model(&self) -> &str {
104        &self.model
105    }
106
107    /// Get the name for this assistant.
108    #[must_use]
109    pub fn name_ref(&self) -> Option<&str> {
110        self.name.as_deref()
111    }
112
113    /// Get the description for this assistant.
114    #[must_use]
115    pub fn description_ref(&self) -> Option<&str> {
116        self.description.as_deref()
117    }
118
119    /// Get the instructions for this assistant.
120    #[must_use]
121    pub fn instructions_ref(&self) -> Option<&str> {
122        self.instructions.as_deref()
123    }
124
125    /// Get the tools for this assistant.
126    #[must_use]
127    pub fn tools_ref(&self) -> &[AssistantTool] {
128        &self.tools
129    }
130
131    /// Get the metadata for this assistant.
132    #[must_use]
133    pub fn metadata_ref(&self) -> &HashMap<String, String> {
134        &self.metadata
135    }
136
137    /// Build a `CreateAssistantRequest` from this builder.
138    pub fn build(self) -> Result<models::CreateAssistantRequest> {
139        let mut request = models::CreateAssistantRequest::new(self.model.clone());
140
141        request.name = self
142            .name
143            .map(|n| Box::new(models::CreateAssistantRequestName::new_text(n)));
144        request.description = self
145            .description
146            .map(|d| Box::new(models::CreateAssistantRequestDescription::new_text(d)));
147        request.instructions = self
148            .instructions
149            .map(|i| Box::new(models::CreateAssistantRequestInstructions::new_text(i)));
150
151        if !self.tools.is_empty() {
152            let tools: Result<Vec<_>> = self
153                .tools
154                .into_iter()
155                .map(|tool| {
156                    match tool {
157                        AssistantTool::CodeInterpreter => Ok(models::AssistantTool::SCode(
158                            Box::new(models::AssistantToolsCode::new(
159                                models::assistant_tools_code::Type::CodeInterpreter,
160                            )),
161                        )),
162                        AssistantTool::FileSearch => Ok(models::AssistantTool::SFileSearch(
163                            Box::new(models::AssistantToolsFileSearch::new(
164                                models::assistant_tools_file_search::Type::FileSearch,
165                            )),
166                        )),
167                        AssistantTool::Function {
168                            name,
169                            description,
170                            parameters,
171                        } => {
172                            let mut function_obj = models::FunctionObject::new(name);
173                            function_obj.description = Some(description);
174                            // Parameters is expected to be a JSON object, so convert it
175                            if let Value::Object(map) = parameters {
176                                let params_map: HashMap<String, Value> = map.into_iter().collect();
177                                function_obj.parameters = Some(params_map);
178                            }
179
180                            let func = models::AssistantToolsFunction::new(
181                                models::assistant_tools_function::Type::Function,
182                                function_obj,
183                            );
184                            Ok(models::AssistantTool::SFunction(Box::new(func)))
185                        }
186                    }
187                })
188                .collect();
189            request.tools = Some(tools?);
190        }
191
192        if !self.metadata.is_empty() {
193            request.metadata = Some(Some(self.metadata.into_iter().collect()));
194        }
195
196        Ok(request)
197    }
198}
199
200/// Represents a tool that can be used by an assistant.
201#[derive(Debug, Clone)]
202pub enum AssistantTool {
203    /// Code interpreter tool for executing Python code.
204    CodeInterpreter,
205    /// File search tool for searching through uploaded files.
206    FileSearch,
207    /// Function calling tool with custom function definition.
208    Function {
209        /// The name of the function.
210        name: String,
211        /// A description of what the function does.
212        description: String,
213        /// The JSON schema that describes the function parameters.
214        parameters: Value,
215    },
216}
217
218/// Builder for creating a thread.
219#[derive(Debug, Clone, Default)]
220pub struct ThreadBuilder {
221    metadata: HashMap<String, String>,
222}
223
224impl ThreadBuilder {
225    /// Create a new thread builder.
226    #[must_use]
227    pub fn new() -> Self {
228        Self::default()
229    }
230
231    /// Add metadata to the thread.
232    #[must_use]
233    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
234        self.metadata.insert(key.into(), value.into());
235        self
236    }
237
238    /// Get the metadata for this thread.
239    #[must_use]
240    pub fn metadata_ref(&self) -> &HashMap<String, String> {
241        &self.metadata
242    }
243}
244
245/// Builder for creating a message.
246#[derive(Debug, Clone)]
247pub struct MessageBuilder {
248    role: String,
249    content: String,
250    attachments: Vec<String>,
251    metadata: HashMap<String, String>,
252}
253
254impl MessageBuilder {
255    /// Create a new message builder with role and content.
256    #[must_use]
257    pub fn new(role: impl Into<String>, content: impl Into<String>) -> Self {
258        Self {
259            role: role.into(),
260            content: content.into(),
261            attachments: Vec::new(),
262            metadata: HashMap::new(),
263        }
264    }
265
266    /// Add an attachment (file ID) to the message.
267    #[must_use]
268    pub fn add_attachment(mut self, file_id: impl Into<String>) -> Self {
269        self.attachments.push(file_id.into());
270        self
271    }
272
273    /// Add metadata to the message.
274    #[must_use]
275    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
276        self.metadata.insert(key.into(), value.into());
277        self
278    }
279
280    /// Get the role for this message.
281    #[must_use]
282    pub fn role_ref(&self) -> &str {
283        &self.role
284    }
285
286    /// Get the content for this message.
287    #[must_use]
288    pub fn content_ref(&self) -> &str {
289        &self.content
290    }
291
292    /// Get the attachments for this message.
293    #[must_use]
294    pub fn attachments_ref(&self) -> &[String] {
295        &self.attachments
296    }
297
298    /// Get the metadata for this message.
299    #[must_use]
300    pub fn metadata_ref(&self) -> &HashMap<String, String> {
301        &self.metadata
302    }
303
304    /// Build a `CreateMessageRequest` from this builder.
305    pub fn build(self) -> Result<models::CreateMessageRequest> {
306        use serde_json::json;
307
308        let role = match self.role.as_str() {
309            "assistant" => models::create_message_request::Role::Assistant,
310            _ => models::create_message_request::Role::User, // Default for "user" and unknown
311        };
312
313        let mut request = models::CreateMessageRequest::new(role, json!(self.content));
314
315        if !self.attachments.is_empty() {
316            let attachments: Vec<_> = self
317                .attachments
318                .into_iter()
319                .map(|file_id| {
320                    let mut att = models::CreateMessageRequestAttachmentsInner::new();
321                    att.file_id = Some(file_id);
322                    att
323                })
324                .collect();
325            request.attachments = Some(Some(attachments));
326        }
327
328        if !self.metadata.is_empty() {
329            request.metadata = Some(Some(self.metadata.into_iter().collect()));
330        }
331
332        Ok(request)
333    }
334}
335
336/// Builder for creating a run.
337#[derive(Debug, Clone)]
338pub struct RunBuilder {
339    assistant_id: String,
340    model: Option<String>,
341    instructions: Option<String>,
342    temperature: Option<f64>,
343    stream: bool,
344    metadata: HashMap<String, String>,
345}
346
347impl RunBuilder {
348    /// Create a new run builder with the specified assistant ID.
349    #[must_use]
350    pub fn new(assistant_id: impl Into<String>) -> Self {
351        Self {
352            assistant_id: assistant_id.into(),
353            model: None,
354            instructions: None,
355            temperature: None,
356            stream: false,
357            metadata: HashMap::new(),
358        }
359    }
360
361    /// Override the assistant's model for this run.
362    #[must_use]
363    pub fn model(mut self, model: impl Into<String>) -> Self {
364        self.model = Some(model.into());
365        self
366    }
367
368    /// Override the assistant's instructions for this run.
369    #[must_use]
370    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
371        self.instructions = Some(instructions.into());
372        self
373    }
374
375    /// Set the temperature for this run.
376    #[must_use]
377    pub fn temperature(mut self, temperature: f64) -> Self {
378        self.temperature = Some(temperature);
379        self
380    }
381
382    /// Enable streaming for this run.
383    #[must_use]
384    pub fn stream(mut self, stream: bool) -> Self {
385        self.stream = stream;
386        self
387    }
388
389    /// Add metadata to the run.
390    #[must_use]
391    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
392        self.metadata.insert(key.into(), value.into());
393        self
394    }
395
396    /// Get the assistant ID for this run.
397    #[must_use]
398    pub fn assistant_id(&self) -> &str {
399        &self.assistant_id
400    }
401
402    /// Get the model override for this run.
403    #[must_use]
404    pub fn model_ref(&self) -> Option<&str> {
405        self.model.as_deref()
406    }
407
408    /// Get the instructions override for this run.
409    #[must_use]
410    pub fn instructions_ref(&self) -> Option<&str> {
411        self.instructions.as_deref()
412    }
413
414    /// Get the temperature for this run.
415    #[must_use]
416    pub fn temperature_ref(&self) -> Option<f64> {
417        self.temperature
418    }
419
420    /// Check if streaming is enabled for this run.
421    #[must_use]
422    pub fn is_streaming(&self) -> bool {
423        self.stream
424    }
425
426    /// Get the metadata for this run.
427    #[must_use]
428    pub fn metadata_ref(&self) -> &HashMap<String, String> {
429        &self.metadata
430    }
431
432    /// Build a `CreateRunRequest` from this builder.
433    pub fn build(self) -> Result<models::CreateRunRequest> {
434        let mut request = models::CreateRunRequest::new(self.assistant_id);
435
436        request.model = self.model;
437        request.instructions = self.instructions;
438        request.temperature = self.temperature;
439        request.stream = Some(self.stream);
440
441        if !self.metadata.is_empty() {
442            request.metadata = Some(Some(self.metadata.into_iter().collect()));
443        }
444
445        Ok(request)
446    }
447}
448
449/// Helper function to create a simple assistant with just a model and name.
450#[must_use]
451pub fn simple_assistant(model: impl Into<String>, name: impl Into<String>) -> AssistantBuilder {
452    AssistantBuilder::new(model).name(name)
453}
454
455/// Helper function to create an assistant with instructions.
456#[must_use]
457pub fn assistant_with_instructions(
458    model: impl Into<String>,
459    name: impl Into<String>,
460    instructions: impl Into<String>,
461) -> AssistantBuilder {
462    AssistantBuilder::new(model)
463        .name(name)
464        .instructions(instructions)
465}
466
467/// Helper function to create a new thread.
468#[must_use]
469pub fn simple_thread() -> ThreadBuilder {
470    ThreadBuilder::new()
471}
472
473/// Helper function to create a simple run.
474#[must_use]
475pub fn simple_run(assistant_id: impl Into<String>) -> RunBuilder {
476    RunBuilder::new(assistant_id)
477}
478
479/// Helper function to create a streaming run.
480#[must_use]
481pub fn streaming_run(assistant_id: impl Into<String>) -> RunBuilder {
482    RunBuilder::new(assistant_id).stream(true)
483}
484
485/// Helper function to create a run with custom temperature.
486#[must_use]
487pub fn temperature_run(assistant_id: impl Into<String>, temperature: f64) -> RunBuilder {
488    RunBuilder::new(assistant_id).temperature(temperature)
489}
490
491/// Helper function to create a code interpreter tool.
492///
493/// The code interpreter tool allows assistants to execute Python code,
494/// perform calculations, data analysis, and generate visualizations.
495///
496/// # Examples
497///
498/// ```rust
499/// use openai_ergonomic::builders::assistants::{AssistantBuilder, tool_code_interpreter};
500///
501/// let assistant = AssistantBuilder::new("gpt-4")
502///     .name("Math Assistant")
503///     .add_tool(tool_code_interpreter());
504/// ```
505#[must_use]
506pub fn tool_code_interpreter() -> AssistantTool {
507    AssistantTool::CodeInterpreter
508}
509
510/// Helper function to create a file search tool.
511///
512/// The file search tool allows assistants to search through uploaded files
513/// and vector stores to provide relevant information from documents.
514///
515/// # Examples
516///
517/// ```rust
518/// use openai_ergonomic::builders::assistants::{AssistantBuilder, tool_file_search};
519///
520/// let assistant = AssistantBuilder::new("gpt-4")
521///     .name("Research Assistant")
522///     .add_tool(tool_file_search());
523/// ```
524#[must_use]
525pub fn tool_file_search() -> AssistantTool {
526    AssistantTool::FileSearch
527}
528
529/// Helper function to create a custom function tool.
530///
531/// Function tools allow assistants to call custom functions that you define,
532/// enabling integration with external APIs and custom business logic.
533///
534/// # Examples
535///
536/// ```rust
537/// use openai_ergonomic::builders::assistants::{AssistantBuilder, tool_function};
538/// use serde_json::json;
539///
540/// let fibonacci_tool = tool_function(
541///     "calculate_fibonacci",
542///     "Calculate the nth Fibonacci number",
543///     json!({
544///         "type": "object",
545///         "properties": {
546///             "n": {
547///                 "type": "integer",
548///                 "description": "The position in the Fibonacci sequence"
549///             }
550///         },
551///         "required": ["n"]
552///     })
553/// );
554///
555/// let assistant = AssistantBuilder::new("gpt-4")
556///     .name("Math Assistant")
557///     .add_tool(fibonacci_tool);
558/// ```
559#[must_use]
560pub fn tool_function(
561    name: impl Into<String>,
562    description: impl Into<String>,
563    parameters: Value,
564) -> AssistantTool {
565    AssistantTool::Function {
566        name: name.into(),
567        description: description.into(),
568        parameters,
569    }
570}
571
572/// Helper function to create an assistant with code interpreter tool.
573#[must_use]
574pub fn assistant_with_code_interpreter(
575    model: impl Into<String>,
576    name: impl Into<String>,
577) -> AssistantBuilder {
578    AssistantBuilder::new(model)
579        .name(name)
580        .add_tool(tool_code_interpreter())
581}
582
583/// Helper function to create an assistant with file search tool.
584#[must_use]
585pub fn assistant_with_file_search(
586    model: impl Into<String>,
587    name: impl Into<String>,
588) -> AssistantBuilder {
589    AssistantBuilder::new(model)
590        .name(name)
591        .add_tool(tool_file_search())
592}
593
594/// Helper function to create an assistant with both code interpreter and file search tools.
595#[must_use]
596pub fn assistant_with_tools(model: impl Into<String>, name: impl Into<String>) -> AssistantBuilder {
597    AssistantBuilder::new(model)
598        .name(name)
599        .add_tool(tool_code_interpreter())
600        .add_tool(tool_file_search())
601}
602
603#[cfg(test)]
604mod tests {
605    use super::*;
606
607    #[test]
608    fn test_assistant_builder() {
609        let builder = AssistantBuilder::new("gpt-4")
610            .name("Test Assistant")
611            .description("A test assistant")
612            .instructions("You are a helpful assistant");
613
614        assert_eq!(builder.model(), "gpt-4");
615        assert_eq!(builder.name_ref(), Some("Test Assistant"));
616        assert_eq!(builder.description_ref(), Some("A test assistant"));
617        assert_eq!(
618            builder.instructions_ref(),
619            Some("You are a helpful assistant")
620        );
621    }
622
623    #[test]
624    fn test_thread_builder() {
625        let builder = ThreadBuilder::new()
626            .metadata("key1", "value1")
627            .metadata("key2", "value2");
628
629        assert_eq!(builder.metadata_ref().len(), 2);
630        assert_eq!(
631            builder.metadata_ref().get("key1"),
632            Some(&"value1".to_string())
633        );
634        assert_eq!(
635            builder.metadata_ref().get("key2"),
636            Some(&"value2".to_string())
637        );
638    }
639
640    #[test]
641    fn test_run_builder() {
642        let builder = RunBuilder::new("assistant-123")
643            .model("gpt-4")
644            .instructions("Follow these instructions")
645            .temperature(0.7)
646            .stream(true)
647            .metadata("key", "value");
648
649        assert_eq!(builder.assistant_id(), "assistant-123");
650        assert_eq!(builder.model_ref(), Some("gpt-4"));
651        assert_eq!(
652            builder.instructions_ref(),
653            Some("Follow these instructions")
654        );
655        assert_eq!(builder.temperature_ref(), Some(0.7));
656        assert!(builder.is_streaming());
657        assert_eq!(builder.metadata_ref().len(), 1);
658    }
659
660    #[test]
661    fn test_simple_assistant_helper() {
662        let builder = simple_assistant("gpt-4", "Helper");
663        assert_eq!(builder.model(), "gpt-4");
664        assert_eq!(builder.name_ref(), Some("Helper"));
665    }
666
667    #[test]
668    fn test_assistant_with_instructions_helper() {
669        let builder = assistant_with_instructions("gpt-4", "Helper", "Be helpful");
670        assert_eq!(builder.model(), "gpt-4");
671        assert_eq!(builder.name_ref(), Some("Helper"));
672        assert_eq!(builder.instructions_ref(), Some("Be helpful"));
673    }
674
675    #[test]
676    fn test_simple_thread_helper() {
677        let builder = simple_thread();
678        assert!(builder.metadata_ref().is_empty());
679    }
680
681    #[test]
682    fn test_simple_run_helper() {
683        let builder = simple_run("assistant-123");
684        assert_eq!(builder.assistant_id(), "assistant-123");
685        assert!(!builder.is_streaming());
686    }
687
688    #[test]
689    fn test_streaming_run_helper() {
690        let builder = streaming_run("assistant-123");
691        assert_eq!(builder.assistant_id(), "assistant-123");
692        assert!(builder.is_streaming());
693    }
694
695    #[test]
696    fn test_temperature_run_helper() {
697        let builder = temperature_run("assistant-123", 0.8);
698        assert_eq!(builder.assistant_id(), "assistant-123");
699        assert_eq!(builder.temperature_ref(), Some(0.8));
700    }
701
702    #[test]
703    fn test_assistant_builder_with_tools() {
704        let builder = AssistantBuilder::new("gpt-4")
705            .name("Tool Assistant")
706            .add_tool(tool_code_interpreter())
707            .add_tool(tool_file_search())
708            .add_metadata("version", "1.0");
709
710        assert_eq!(builder.model(), "gpt-4");
711        assert_eq!(builder.name_ref(), Some("Tool Assistant"));
712        assert_eq!(builder.tools_ref().len(), 2);
713        assert_eq!(builder.metadata_ref().len(), 1);
714
715        // Check tool types
716        match &builder.tools_ref()[0] {
717            AssistantTool::CodeInterpreter => {}
718            _ => panic!("Expected CodeInterpreter tool"),
719        }
720
721        match &builder.tools_ref()[1] {
722            AssistantTool::FileSearch => {}
723            _ => panic!("Expected FileSearch tool"),
724        }
725    }
726
727    #[test]
728    fn test_tool_function() {
729        use serde_json::json;
730
731        let tool = tool_function(
732            "test_function",
733            "A test function",
734            json!({"type": "object", "properties": {}}),
735        );
736
737        match tool {
738            AssistantTool::Function {
739                name,
740                description,
741                parameters,
742            } => {
743                assert_eq!(name, "test_function");
744                assert_eq!(description, "A test function");
745                assert!(parameters.is_object());
746            }
747            _ => panic!("Expected Function tool"),
748        }
749    }
750
751    #[test]
752    fn test_tool_helpers() {
753        let code_tool = tool_code_interpreter();
754        match code_tool {
755            AssistantTool::CodeInterpreter => {}
756            _ => panic!("Expected CodeInterpreter tool"),
757        }
758
759        let search_tool = tool_file_search();
760        match search_tool {
761            AssistantTool::FileSearch => {}
762            _ => panic!("Expected FileSearch tool"),
763        }
764    }
765
766    #[test]
767    fn test_assistant_with_code_interpreter_helper() {
768        let builder = assistant_with_code_interpreter("gpt-4", "Code Helper");
769        assert_eq!(builder.model(), "gpt-4");
770        assert_eq!(builder.name_ref(), Some("Code Helper"));
771        assert_eq!(builder.tools_ref().len(), 1);
772
773        match &builder.tools_ref()[0] {
774            AssistantTool::CodeInterpreter => {}
775            _ => panic!("Expected CodeInterpreter tool"),
776        }
777    }
778
779    #[test]
780    fn test_assistant_with_file_search_helper() {
781        let builder = assistant_with_file_search("gpt-4", "Search Helper");
782        assert_eq!(builder.model(), "gpt-4");
783        assert_eq!(builder.name_ref(), Some("Search Helper"));
784        assert_eq!(builder.tools_ref().len(), 1);
785
786        match &builder.tools_ref()[0] {
787            AssistantTool::FileSearch => {}
788            _ => panic!("Expected FileSearch tool"),
789        }
790    }
791
792    #[test]
793    fn test_assistant_with_tools_helper() {
794        let builder = assistant_with_tools("gpt-4", "Multi-Tool Helper");
795        assert_eq!(builder.model(), "gpt-4");
796        assert_eq!(builder.name_ref(), Some("Multi-Tool Helper"));
797        assert_eq!(builder.tools_ref().len(), 2);
798
799        match &builder.tools_ref()[0] {
800            AssistantTool::CodeInterpreter => {}
801            _ => panic!("Expected CodeInterpreter tool"),
802        }
803
804        match &builder.tools_ref()[1] {
805            AssistantTool::FileSearch => {}
806            _ => panic!("Expected FileSearch tool"),
807        }
808    }
809}