openai_ergonomic/builders/
chat.rs

1//! Chat completion builders and helpers.
2//!
3//! This module provides ergonomic builders for chat completion requests,
4//! including helpers for common message patterns and streaming responses.
5
6use openai_client_base::models::{
7    ChatCompletionRequestAssistantMessage, ChatCompletionRequestAssistantMessageContent,
8    ChatCompletionRequestMessage, ChatCompletionRequestMessageContentPartImage,
9    ChatCompletionRequestMessageContentPartImageImageUrl,
10    ChatCompletionRequestMessageContentPartText, ChatCompletionRequestSystemMessage,
11    ChatCompletionRequestSystemMessageContent, ChatCompletionRequestToolMessage,
12    ChatCompletionRequestToolMessageContent, ChatCompletionRequestUserMessage,
13    ChatCompletionRequestUserMessageContent, ChatCompletionRequestUserMessageContentPart,
14    ChatCompletionTool, ChatCompletionToolChoiceOption, CreateChatCompletionRequest,
15    CreateChatCompletionRequestAllOfTools, FunctionObject,
16};
17// Import the specific Role enums for each message type
18use openai_client_base::models::chat_completion_request_assistant_message::Role as AssistantRole;
19use openai_client_base::models::chat_completion_request_system_message::Role as SystemRole;
20use openai_client_base::models::chat_completion_request_tool_message::Role as ToolRole;
21use openai_client_base::models::chat_completion_request_user_message::Role as UserRole;
22// Import the Type enums for content parts
23use openai_client_base::models::chat_completion_request_message_content_part_image::Type as ImageType;
24use openai_client_base::models::chat_completion_request_message_content_part_image_image_url::Detail;
25use openai_client_base::models::chat_completion_request_message_content_part_text::Type as TextType;
26use serde_json::Value;
27
28/// Builder for chat completion requests.
29#[derive(Debug, Clone)]
30pub struct ChatCompletionBuilder {
31    model: String,
32    messages: Vec<ChatCompletionRequestMessage>,
33    temperature: Option<f64>,
34    max_tokens: Option<i32>,
35    max_completion_tokens: Option<i32>,
36    stream: Option<bool>,
37    tools: Option<Vec<ChatCompletionTool>>,
38    tool_choice: Option<ChatCompletionToolChoiceOption>,
39    response_format:
40        Option<openai_client_base::models::CreateChatCompletionRequestAllOfResponseFormat>,
41    n: Option<i32>,
42    stop: Option<Vec<String>>,
43    presence_penalty: Option<f64>,
44    frequency_penalty: Option<f64>,
45    top_p: Option<f64>,
46    user: Option<String>,
47    seed: Option<i32>,
48}
49
50impl ChatCompletionBuilder {
51    /// Create a new chat completion builder with the specified model.
52    #[must_use]
53    pub fn new(model: impl Into<String>) -> Self {
54        Self {
55            model: model.into(),
56            messages: Vec::new(),
57            temperature: None,
58            max_tokens: None,
59            max_completion_tokens: None,
60            stream: None,
61            tools: None,
62            tool_choice: None,
63            response_format: None,
64            n: None,
65            stop: None,
66            presence_penalty: None,
67            frequency_penalty: None,
68            top_p: None,
69            user: None,
70            seed: None,
71        }
72    }
73
74    /// Add a system message to the conversation.
75    #[must_use]
76    pub fn system(mut self, content: impl Into<String>) -> Self {
77        let message = ChatCompletionRequestSystemMessage {
78            content: Box::new(ChatCompletionRequestSystemMessageContent::TextContent(
79                content.into(),
80            )),
81            role: SystemRole::System,
82            name: None,
83        };
84        self.messages.push(
85            ChatCompletionRequestMessage::ChatCompletionRequestSystemMessage(Box::new(message)),
86        );
87        self
88    }
89
90    /// Add a user message to the conversation.
91    #[must_use]
92    pub fn user(mut self, content: impl Into<String>) -> Self {
93        let message = ChatCompletionRequestUserMessage {
94            content: Box::new(ChatCompletionRequestUserMessageContent::TextContent(
95                content.into(),
96            )),
97            role: UserRole::User,
98            name: None,
99        };
100        self.messages.push(
101            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(Box::new(message)),
102        );
103        self
104    }
105
106    /// Add a user message with both text and an image URL.
107    #[must_use]
108    pub fn user_with_image_url(
109        self,
110        text: impl Into<String>,
111        image_url: impl Into<String>,
112    ) -> Self {
113        self.user_with_image_url_and_detail(text, image_url, Detail::Auto)
114    }
115
116    /// Add a user message with both text and an image URL with specified detail level.
117    #[must_use]
118    pub fn user_with_image_url_and_detail(
119        mut self,
120        text: impl Into<String>,
121        image_url: impl Into<String>,
122        detail: Detail,
123    ) -> Self {
124        let text_part = ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartText(
125            Box::new(ChatCompletionRequestMessageContentPartText {
126                r#type: TextType::Text,
127                text: text.into(),
128            }),
129        );
130
131        let image_part = ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(
132            Box::new(ChatCompletionRequestMessageContentPartImage {
133                r#type: ImageType::ImageUrl,
134                image_url: Box::new(ChatCompletionRequestMessageContentPartImageImageUrl {
135                    url: image_url.into(),
136                    detail: Some(detail),
137                }),
138            }),
139        );
140
141        let message = ChatCompletionRequestUserMessage {
142            content: Box::new(
143                ChatCompletionRequestUserMessageContent::ArrayOfContentParts(vec![
144                    text_part, image_part,
145                ]),
146            ),
147            role: UserRole::User,
148            name: None,
149        };
150
151        self.messages.push(
152            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(Box::new(message)),
153        );
154        self
155    }
156
157    /// Add a user message with multiple content parts (text and/or images).
158    #[must_use]
159    pub fn user_with_parts(
160        mut self,
161        parts: Vec<ChatCompletionRequestUserMessageContentPart>,
162    ) -> Self {
163        let message = ChatCompletionRequestUserMessage {
164            content: Box::new(ChatCompletionRequestUserMessageContent::ArrayOfContentParts(parts)),
165            role: UserRole::User,
166            name: None,
167        };
168
169        self.messages.push(
170            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(Box::new(message)),
171        );
172        self
173    }
174
175    /// Add an assistant message to the conversation.
176    #[must_use]
177    pub fn assistant(mut self, content: impl Into<String>) -> Self {
178        let message = ChatCompletionRequestAssistantMessage {
179            content: Some(Some(Box::new(
180                ChatCompletionRequestAssistantMessageContent::TextContent(content.into()),
181            ))),
182            role: AssistantRole::Assistant,
183            name: None,
184            tool_calls: None,
185            function_call: None,
186            audio: None,
187            refusal: None,
188        };
189        self.messages.push(
190            ChatCompletionRequestMessage::ChatCompletionRequestAssistantMessage(Box::new(message)),
191        );
192        self
193    }
194
195    /// Add an assistant message with tool calls to the conversation.
196    ///
197    /// This is used when the assistant wants to call tools. Each tool call
198    /// should be represented as a tuple of (`tool_call_id`, `function_name`, `function_arguments`).
199    #[must_use]
200    pub fn assistant_with_tool_calls(
201        mut self,
202        content: impl Into<String>,
203        tool_calls: Vec<openai_client_base::models::ChatCompletionMessageToolCallsInner>,
204    ) -> Self {
205        let content_str = content.into();
206        let message = ChatCompletionRequestAssistantMessage {
207            content: if content_str.is_empty() {
208                None
209            } else {
210                Some(Some(Box::new(
211                    ChatCompletionRequestAssistantMessageContent::TextContent(content_str),
212                )))
213            },
214            role: AssistantRole::Assistant,
215            name: None,
216            tool_calls: Some(tool_calls),
217            function_call: None,
218            audio: None,
219            refusal: None,
220        };
221        self.messages.push(
222            ChatCompletionRequestMessage::ChatCompletionRequestAssistantMessage(Box::new(message)),
223        );
224        self
225    }
226
227    /// Add a tool result message to the conversation.
228    ///
229    /// This is used to provide the result of a tool call back to the assistant.
230    #[must_use]
231    pub fn tool(mut self, tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
232        let message = ChatCompletionRequestToolMessage {
233            role: ToolRole::Tool,
234            content: Box::new(ChatCompletionRequestToolMessageContent::TextContent(
235                content.into(),
236            )),
237            tool_call_id: tool_call_id.into(),
238        };
239        self.messages.push(
240            ChatCompletionRequestMessage::ChatCompletionRequestToolMessage(Box::new(message)),
241        );
242        self
243    }
244
245    /// Set the temperature for the completion.
246    #[must_use]
247    pub fn temperature(mut self, temperature: f64) -> Self {
248        self.temperature = Some(temperature);
249        self
250    }
251
252    /// Set the maximum number of tokens to generate.
253    #[must_use]
254    pub fn max_tokens(mut self, max_tokens: i32) -> Self {
255        self.max_tokens = Some(max_tokens);
256        self
257    }
258
259    /// Set the maximum completion tokens (for newer models).
260    #[must_use]
261    pub fn max_completion_tokens(mut self, max_completion_tokens: i32) -> Self {
262        self.max_completion_tokens = Some(max_completion_tokens);
263        self
264    }
265
266    /// Enable streaming for the completion.
267    #[must_use]
268    pub fn stream(mut self, stream: bool) -> Self {
269        self.stream = Some(stream);
270        self
271    }
272
273    /// Add tools that the model can use.
274    #[must_use]
275    pub fn tools(mut self, tools: Vec<ChatCompletionTool>) -> Self {
276        self.tools = Some(tools);
277        self
278    }
279
280    /// Set the tool choice option.
281    #[must_use]
282    pub fn tool_choice(mut self, tool_choice: ChatCompletionToolChoiceOption) -> Self {
283        self.tool_choice = Some(tool_choice);
284        self
285    }
286
287    /// Set the response format.
288    #[must_use]
289    pub fn response_format(
290        mut self,
291        format: openai_client_base::models::CreateChatCompletionRequestAllOfResponseFormat,
292    ) -> Self {
293        self.response_format = Some(format);
294        self
295    }
296
297    /// Set the number of completions to generate.
298    #[must_use]
299    pub fn n(mut self, n: i32) -> Self {
300        self.n = Some(n);
301        self
302    }
303
304    /// Set stop sequences.
305    #[must_use]
306    pub fn stop(mut self, stop: Vec<String>) -> Self {
307        self.stop = Some(stop);
308        self
309    }
310
311    /// Set the presence penalty.
312    #[must_use]
313    pub fn presence_penalty(mut self, presence_penalty: f64) -> Self {
314        self.presence_penalty = Some(presence_penalty);
315        self
316    }
317
318    /// Set the frequency penalty.
319    #[must_use]
320    pub fn frequency_penalty(mut self, frequency_penalty: f64) -> Self {
321        self.frequency_penalty = Some(frequency_penalty);
322        self
323    }
324
325    /// Set the top-p value.
326    #[must_use]
327    pub fn top_p(mut self, top_p: f64) -> Self {
328        self.top_p = Some(top_p);
329        self
330    }
331
332    /// Set the user identifier.
333    #[must_use]
334    pub fn user_id(mut self, user: impl Into<String>) -> Self {
335        self.user = Some(user.into());
336        self
337    }
338
339    /// Set the random seed for deterministic outputs.
340    #[must_use]
341    pub fn seed(mut self, seed: i32) -> Self {
342        self.seed = Some(seed);
343        self
344    }
345}
346
347impl super::Builder<CreateChatCompletionRequest> for ChatCompletionBuilder {
348    #[allow(clippy::too_many_lines)]
349    fn build(self) -> crate::Result<CreateChatCompletionRequest> {
350        // Validate model
351        if self.model.trim().is_empty() {
352            return Err(crate::Error::InvalidRequest(
353                "Model cannot be empty".to_string(),
354            ));
355        }
356
357        // Validate messages
358        if self.messages.is_empty() {
359            return Err(crate::Error::InvalidRequest(
360                "At least one message is required".to_string(),
361            ));
362        }
363
364        // Validate message contents
365        for (i, message) in self.messages.iter().enumerate() {
366            match message {
367                ChatCompletionRequestMessage::ChatCompletionRequestSystemMessage(msg) => {
368                    if let ChatCompletionRequestSystemMessageContent::TextContent(content) =
369                        msg.content.as_ref()
370                    {
371                        if content.trim().is_empty() {
372                            return Err(crate::Error::InvalidRequest(format!(
373                                "System message at index {i} cannot have empty content"
374                            )));
375                        }
376                    }
377                }
378                ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(msg) => {
379                    match msg.content.as_ref() {
380                        ChatCompletionRequestUserMessageContent::TextContent(content) => {
381                            if content.trim().is_empty() {
382                                return Err(crate::Error::InvalidRequest(format!(
383                                    "User message at index {i} cannot have empty content"
384                                )));
385                            }
386                        }
387                        ChatCompletionRequestUserMessageContent::ArrayOfContentParts(parts) => {
388                            if parts.is_empty() {
389                                return Err(crate::Error::InvalidRequest(format!(
390                                    "User message at index {i} cannot have empty content parts"
391                                )));
392                            }
393                        }
394                    }
395                }
396                ChatCompletionRequestMessage::ChatCompletionRequestAssistantMessage(msg) => {
397                    // Assistant messages can have content or tool calls, but not both empty
398                    let has_content = msg
399                        .content
400                        .as_ref()
401                        .and_then(|opt| opt.as_ref())
402                        .is_some_and(|c| {
403                            match c.as_ref() {
404                                ChatCompletionRequestAssistantMessageContent::TextContent(text) => {
405                                    !text.trim().is_empty()
406                                }
407                                _ => true, // Other content types are considered valid
408                            }
409                        });
410                    let has_tool_calls = msg.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty());
411
412                    if !has_content && !has_tool_calls {
413                        return Err(crate::Error::InvalidRequest(format!(
414                            "Assistant message at index {i} must have either content or tool calls"
415                        )));
416                    }
417                }
418                _ => {
419                    // Other message types (tool, function) are valid as-is
420                }
421            }
422        }
423
424        // Validate temperature
425        if let Some(temp) = self.temperature {
426            if !(0.0..=2.0).contains(&temp) {
427                return Err(crate::Error::InvalidRequest(format!(
428                    "temperature must be between 0.0 and 2.0, got {temp}"
429                )));
430            }
431        }
432
433        // Validate top_p
434        if let Some(top_p) = self.top_p {
435            if !(0.0..=1.0).contains(&top_p) {
436                return Err(crate::Error::InvalidRequest(format!(
437                    "top_p must be between 0.0 and 1.0, got {top_p}"
438                )));
439            }
440        }
441
442        // Validate frequency_penalty
443        if let Some(freq) = self.frequency_penalty {
444            if !(-2.0..=2.0).contains(&freq) {
445                return Err(crate::Error::InvalidRequest(format!(
446                    "frequency_penalty must be between -2.0 and 2.0, got {freq}"
447                )));
448            }
449        }
450
451        // Validate presence_penalty
452        if let Some(pres) = self.presence_penalty {
453            if !(-2.0..=2.0).contains(&pres) {
454                return Err(crate::Error::InvalidRequest(format!(
455                    "presence_penalty must be between -2.0 and 2.0, got {pres}"
456                )));
457            }
458        }
459
460        // Validate max_tokens
461        if let Some(max_tokens) = self.max_tokens {
462            if max_tokens <= 0 {
463                return Err(crate::Error::InvalidRequest(format!(
464                    "max_tokens must be positive, got {max_tokens}"
465                )));
466            }
467        }
468
469        // Validate max_completion_tokens
470        if let Some(max_completion_tokens) = self.max_completion_tokens {
471            if max_completion_tokens <= 0 {
472                return Err(crate::Error::InvalidRequest(format!(
473                    "max_completion_tokens must be positive, got {max_completion_tokens}"
474                )));
475            }
476        }
477
478        // Validate n
479        if let Some(n) = self.n {
480            if n <= 0 {
481                return Err(crate::Error::InvalidRequest(format!(
482                    "n must be positive, got {n}"
483                )));
484            }
485        }
486
487        // Validate tools
488        if let Some(ref tools) = self.tools {
489            for (i, tool) in tools.iter().enumerate() {
490                let function = &tool.function;
491
492                // Validate function name
493                if function.name.trim().is_empty() {
494                    return Err(crate::Error::InvalidRequest(format!(
495                        "Tool {i} function name cannot be empty"
496                    )));
497                }
498
499                // Validate function name contains only valid characters
500                if !function
501                    .name
502                    .chars()
503                    .all(|c| c.is_alphanumeric() || c == '_')
504                {
505                    return Err(crate::Error::InvalidRequest(format!(
506                        "Tool {} function name '{}' contains invalid characters",
507                        i, function.name
508                    )));
509                }
510
511                // Validate function description
512                if let Some(ref description) = &function.description {
513                    if description.trim().is_empty() {
514                        return Err(crate::Error::InvalidRequest(format!(
515                            "Tool {i} function description cannot be empty"
516                        )));
517                    }
518                }
519            }
520        }
521
522        let response_format = self.response_format.map(Box::new);
523
524        Ok(CreateChatCompletionRequest {
525            messages: self.messages,
526            model: self.model,
527            frequency_penalty: self.frequency_penalty,
528            logit_bias: None,
529            logprobs: None,
530            top_logprobs: None,
531            max_tokens: self.max_tokens,
532            max_completion_tokens: self.max_completion_tokens,
533            n: self.n,
534            modalities: None,
535            prediction: None,
536            audio: None,
537            presence_penalty: self.presence_penalty,
538            response_format,
539            seed: self.seed,
540            service_tier: None,
541            stop: self.stop.map(|s| {
542                Box::new(openai_client_base::models::StopConfiguration::ArrayOfStrings(s))
543            }),
544            stream: self.stream,
545            stream_options: None,
546            temperature: self.temperature,
547            top_p: self.top_p,
548            tools: self.tools.map(|tools| {
549                tools
550                    .into_iter()
551                    .map(|tool| {
552                        CreateChatCompletionRequestAllOfTools::ChatCompletionTool(Box::new(tool))
553                    })
554                    .collect()
555            }),
556            tool_choice: self.tool_choice.map(Box::new),
557            parallel_tool_calls: None,
558            user: self.user,
559            function_call: None,
560            functions: None,
561            store: None,
562            metadata: None,
563            reasoning_effort: None,
564            prompt_cache_key: None,
565            safety_identifier: None,
566            verbosity: None,
567            web_search_options: None,
568        })
569    }
570}
571
572// TODO: Implement Sendable trait once client is available
573// impl super::Sendable<ChatCompletionResponse> for ChatCompletionBuilder {
574//     async fn send(self) -> crate::Result<ChatCompletionResponse> {
575//         // Implementation will use the client wrapper
576//         todo!("Implement once client wrapper is available")
577//     }
578// }
579
580/// Helper function to create a simple user message chat completion.
581#[must_use]
582pub fn user_message(model: impl Into<String>, content: impl Into<String>) -> ChatCompletionBuilder {
583    ChatCompletionBuilder::new(model).user(content)
584}
585
586/// Helper function to create a system + user message chat completion.
587#[must_use]
588pub fn system_user(
589    model: impl Into<String>,
590    system: impl Into<String>,
591    user: impl Into<String>,
592) -> ChatCompletionBuilder {
593    ChatCompletionBuilder::new(model).system(system).user(user)
594}
595
596/// Helper function to create a function tool.
597#[must_use]
598pub fn tool_function(
599    name: impl Into<String>,
600    description: impl Into<String>,
601    parameters: Value,
602) -> ChatCompletionTool {
603    use std::collections::HashMap;
604
605    // Convert Value to HashMap<String, Value>
606    let params_map = if let serde_json::Value::Object(map) = parameters {
607        map.into_iter().collect::<HashMap<String, Value>>()
608    } else {
609        HashMap::new()
610    };
611
612    ChatCompletionTool {
613        r#type: openai_client_base::models::chat_completion_tool::Type::Function,
614        function: Box::new(FunctionObject {
615            name: name.into(),
616            description: Some(description.into()),
617            parameters: Some(params_map),
618            strict: None,
619        }),
620    }
621}
622
623/// Helper function to create a web search tool.
624#[must_use]
625pub fn tool_web_search() -> ChatCompletionTool {
626    tool_function(
627        "web_search",
628        "Search the web for information",
629        serde_json::json!({
630            "type": "object",
631            "properties": {
632                "query": {
633                    "type": "string",
634                    "description": "The search query"
635                }
636            },
637            "required": ["query"]
638        }),
639    )
640}
641
642/// Helper function to create a text content part.
643#[must_use]
644pub fn text_part(content: impl Into<String>) -> ChatCompletionRequestUserMessageContentPart {
645    ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartText(
646        Box::new(ChatCompletionRequestMessageContentPartText {
647            r#type: TextType::Text,
648            text: content.into(),
649        }),
650    )
651}
652
653/// Helper function to create an image content part from a URL with auto detail.
654#[must_use]
655pub fn image_url_part(url: impl Into<String>) -> ChatCompletionRequestUserMessageContentPart {
656    image_url_part_with_detail(url, Detail::Auto)
657}
658
659/// Helper function to create an image content part from a URL with specified detail level.
660#[must_use]
661pub fn image_url_part_with_detail(
662    url: impl Into<String>,
663    detail: Detail,
664) -> ChatCompletionRequestUserMessageContentPart {
665    ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(
666        Box::new(ChatCompletionRequestMessageContentPartImage {
667            r#type: ImageType::ImageUrl,
668            image_url: Box::new(ChatCompletionRequestMessageContentPartImageImageUrl {
669                url: url.into(),
670                detail: Some(detail),
671            }),
672        }),
673    )
674}
675
676/// Helper function to create an image content part from base64 data with auto detail.
677#[must_use]
678pub fn image_base64_part(
679    base64_data: impl Into<String>,
680    media_type: impl Into<String>,
681) -> ChatCompletionRequestUserMessageContentPart {
682    image_base64_part_with_detail(base64_data, media_type, Detail::Auto)
683}
684
685/// Helper function to create an image content part from base64 data with specified detail level.
686#[must_use]
687pub fn image_base64_part_with_detail(
688    base64_data: impl Into<String>,
689    media_type: impl Into<String>,
690    detail: Detail,
691) -> ChatCompletionRequestUserMessageContentPart {
692    let data_url = format!("data:{};base64,{}", media_type.into(), base64_data.into());
693    image_url_part_with_detail(data_url, detail)
694}
695
696#[cfg(test)]
697mod tests {
698    use super::*;
699    use crate::builders::Builder;
700    use openai_client_base::models::chat_completion_tool_choice_option::ChatCompletionToolChoiceOption;
701
702    #[test]
703    fn test_chat_completion_builder_new() {
704        let builder = ChatCompletionBuilder::new("gpt-4");
705        assert_eq!(builder.model, "gpt-4");
706        assert!(builder.messages.is_empty());
707        assert!(builder.temperature.is_none());
708    }
709
710    #[test]
711    fn test_chat_completion_builder_system_message() {
712        let builder = ChatCompletionBuilder::new("gpt-4").system("You are a helpful assistant");
713        assert_eq!(builder.messages.len(), 1);
714
715        // Verify the message structure
716        match &builder.messages[0] {
717            ChatCompletionRequestMessage::ChatCompletionRequestSystemMessage(msg) => {
718                match msg.content.as_ref() {
719                    ChatCompletionRequestSystemMessageContent::TextContent(content) => {
720                        assert_eq!(content, "You are a helpful assistant");
721                    }
722                    ChatCompletionRequestSystemMessageContent::ArrayOfContentParts(_) => {
723                        panic!("Expected text content")
724                    }
725                }
726            }
727            _ => panic!("Expected system message"),
728        }
729    }
730
731    #[test]
732    fn test_chat_completion_builder_user_message() {
733        let builder = ChatCompletionBuilder::new("gpt-4").user("Hello, world!");
734        assert_eq!(builder.messages.len(), 1);
735
736        // Verify the message structure
737        match &builder.messages[0] {
738            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(msg) => {
739                match msg.content.as_ref() {
740                    ChatCompletionRequestUserMessageContent::TextContent(content) => {
741                        assert_eq!(content, "Hello, world!");
742                    }
743                    ChatCompletionRequestUserMessageContent::ArrayOfContentParts(_) => {
744                        panic!("Expected text content")
745                    }
746                }
747            }
748            _ => panic!("Expected user message"),
749        }
750    }
751
752    #[test]
753    fn test_chat_completion_builder_assistant_message() {
754        let builder = ChatCompletionBuilder::new("gpt-4").assistant("Hello! How can I help you?");
755        assert_eq!(builder.messages.len(), 1);
756
757        // Verify the message structure
758        match &builder.messages[0] {
759            ChatCompletionRequestMessage::ChatCompletionRequestAssistantMessage(msg) => {
760                if let Some(Some(content)) = &msg.content {
761                    match content.as_ref() {
762                        ChatCompletionRequestAssistantMessageContent::TextContent(text) => {
763                            assert_eq!(text, "Hello! How can I help you?");
764                        }
765                        _ => panic!("Expected text content"),
766                    }
767                } else {
768                    panic!("Expected content");
769                }
770            }
771            _ => panic!("Expected assistant message"),
772        }
773    }
774
775    #[test]
776    fn test_chat_completion_builder_user_with_image_url() {
777        let builder = ChatCompletionBuilder::new("gpt-4")
778            .user_with_image_url("Describe this image", "https://example.com/image.jpg");
779        assert_eq!(builder.messages.len(), 1);
780
781        // Verify the message structure
782        match &builder.messages[0] {
783            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(msg) => {
784                match msg.content.as_ref() {
785                    ChatCompletionRequestUserMessageContent::ArrayOfContentParts(parts) => {
786                        assert_eq!(parts.len(), 2);
787
788                        // Check text part
789                        match &parts[0] {
790                            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartText(text_part) => {
791                                assert_eq!(text_part.text, "Describe this image");
792                            }
793                            _ => panic!("Expected text part"),
794                        }
795
796                        // Check image part
797                        match &parts[1] {
798                            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(image_part) => {
799                                assert_eq!(image_part.image_url.url, "https://example.com/image.jpg");
800                                assert_eq!(image_part.image_url.detail, Some(Detail::Auto));
801                            }
802                            _ => panic!("Expected image part"),
803                        }
804                    }
805                    ChatCompletionRequestUserMessageContent::TextContent(_) => {
806                        panic!("Expected array of content parts")
807                    }
808                }
809            }
810            _ => panic!("Expected user message"),
811        }
812    }
813
814    #[test]
815    fn test_chat_completion_builder_user_with_image_url_and_detail() {
816        let builder = ChatCompletionBuilder::new("gpt-4").user_with_image_url_and_detail(
817            "Describe this image",
818            "https://example.com/image.jpg",
819            Detail::High,
820        );
821        assert_eq!(builder.messages.len(), 1);
822
823        // Verify the message structure
824        match &builder.messages[0] {
825            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(msg) => {
826                match msg.content.as_ref() {
827                    ChatCompletionRequestUserMessageContent::ArrayOfContentParts(parts) => {
828                        assert_eq!(parts.len(), 2);
829
830                        // Check image part detail
831                        match &parts[1] {
832                            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(image_part) => {
833                                assert_eq!(image_part.image_url.detail, Some(Detail::High));
834                            }
835                            _ => panic!("Expected image part"),
836                        }
837                    }
838                    ChatCompletionRequestUserMessageContent::TextContent(_) => {
839                        panic!("Expected array of content parts")
840                    }
841                }
842            }
843            _ => panic!("Expected user message"),
844        }
845    }
846
847    #[test]
848    fn test_chat_completion_builder_user_with_parts() {
849        let text_part = text_part("Hello");
850        let image_part = image_url_part("https://example.com/image.jpg");
851        let parts = vec![text_part, image_part];
852
853        let builder = ChatCompletionBuilder::new("gpt-4").user_with_parts(parts);
854        assert_eq!(builder.messages.len(), 1);
855
856        // Verify the message structure
857        match &builder.messages[0] {
858            ChatCompletionRequestMessage::ChatCompletionRequestUserMessage(msg) => {
859                match msg.content.as_ref() {
860                    ChatCompletionRequestUserMessageContent::ArrayOfContentParts(parts) => {
861                        assert_eq!(parts.len(), 2);
862                    }
863                    ChatCompletionRequestUserMessageContent::TextContent(_) => {
864                        panic!("Expected array of content parts")
865                    }
866                }
867            }
868            _ => panic!("Expected user message"),
869        }
870    }
871
872    #[test]
873    fn test_chat_completion_builder_chaining() {
874        let builder = ChatCompletionBuilder::new("gpt-4")
875            .system("You are a helpful assistant")
876            .user("What's the weather?")
877            .temperature(0.7)
878            .max_tokens(100);
879
880        assert_eq!(builder.messages.len(), 2);
881        assert_eq!(builder.temperature, Some(0.7));
882        assert_eq!(builder.max_tokens, Some(100));
883    }
884
885    #[test]
886    fn test_chat_completion_builder_parameters() {
887        let builder = ChatCompletionBuilder::new("gpt-4")
888            .temperature(0.5)
889            .max_tokens(150)
890            .max_completion_tokens(200)
891            .stream(true)
892            .n(2)
893            .stop(vec!["STOP".to_string()])
894            .presence_penalty(0.1)
895            .frequency_penalty(0.2)
896            .top_p(0.9)
897            .user_id("user123");
898
899        assert_eq!(builder.temperature, Some(0.5));
900        assert_eq!(builder.max_tokens, Some(150));
901        assert_eq!(builder.max_completion_tokens, Some(200));
902        assert_eq!(builder.stream, Some(true));
903        assert_eq!(builder.n, Some(2));
904        assert_eq!(builder.stop, Some(vec!["STOP".to_string()]));
905        assert_eq!(builder.presence_penalty, Some(0.1));
906        assert_eq!(builder.frequency_penalty, Some(0.2));
907        assert_eq!(builder.top_p, Some(0.9));
908        assert_eq!(builder.user, Some("user123".to_string()));
909    }
910
911    #[test]
912    fn test_chat_completion_builder_tools() {
913        let tool = tool_function(
914            "test_function",
915            "A test function",
916            serde_json::json!({"type": "object", "properties": {}}),
917        );
918
919        let builder = ChatCompletionBuilder::new("gpt-4")
920            .tools(vec![tool])
921            .tool_choice(ChatCompletionToolChoiceOption::Auto(
922                openai_client_base::models::chat_completion_tool_choice_option::ChatCompletionToolChoiceOptionAutoEnum::Auto
923            ));
924
925        assert_eq!(builder.tools.as_ref().unwrap().len(), 1);
926        assert!(builder.tool_choice.is_some());
927    }
928
929    #[test]
930    fn test_chat_completion_builder_build_success() {
931        let builder = ChatCompletionBuilder::new("gpt-4").user("Hello");
932        let request = builder.build().unwrap();
933
934        assert_eq!(request.model, "gpt-4");
935        assert_eq!(request.messages.len(), 1);
936    }
937
938    #[test]
939    fn test_chat_completion_builder_build_empty_messages_error() {
940        let builder = ChatCompletionBuilder::new("gpt-4");
941        let result = builder.build();
942
943        assert!(result.is_err());
944        if let Err(error) = result {
945            assert!(matches!(error, crate::Error::InvalidRequest(_)));
946        }
947    }
948
949    #[test]
950    fn test_user_message_helper() {
951        let builder = user_message("gpt-4", "Hello, world!");
952        assert_eq!(builder.model, "gpt-4");
953        assert_eq!(builder.messages.len(), 1);
954    }
955
956    #[test]
957    fn test_system_user_helper() {
958        let builder = system_user(
959            "gpt-4",
960            "You are a helpful assistant",
961            "What's the weather?",
962        );
963        assert_eq!(builder.model, "gpt-4");
964        assert_eq!(builder.messages.len(), 2);
965    }
966
967    #[test]
968    fn test_tool_function() {
969        let tool = tool_function(
970            "get_weather",
971            "Get current weather",
972            serde_json::json!({
973                "type": "object",
974                "properties": {
975                    "location": {"type": "string"}
976                }
977            }),
978        );
979
980        assert_eq!(tool.function.name, "get_weather");
981        assert_eq!(
982            tool.function.description.as_ref().unwrap(),
983            "Get current weather"
984        );
985        assert!(tool.function.parameters.is_some());
986    }
987
988    #[test]
989    fn test_tool_web_search() {
990        let tool = tool_web_search();
991        assert_eq!(tool.function.name, "web_search");
992        assert!(tool.function.description.is_some());
993        assert!(tool.function.parameters.is_some());
994    }
995
996    #[test]
997    fn test_text_part() {
998        let part = text_part("Hello, world!");
999        match part {
1000            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartText(text_part) => {
1001                assert_eq!(text_part.text, "Hello, world!");
1002                assert_eq!(text_part.r#type, TextType::Text);
1003            }
1004            _ => panic!("Expected text part"),
1005        }
1006    }
1007
1008    #[test]
1009    fn test_image_url_part() {
1010        let part = image_url_part("https://example.com/image.jpg");
1011        match part {
1012            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(image_part) => {
1013                assert_eq!(image_part.image_url.url, "https://example.com/image.jpg");
1014                assert_eq!(image_part.image_url.detail, Some(Detail::Auto));
1015                assert_eq!(image_part.r#type, ImageType::ImageUrl);
1016            }
1017            _ => panic!("Expected image part"),
1018        }
1019    }
1020
1021    #[test]
1022    fn test_image_url_part_with_detail() {
1023        let part = image_url_part_with_detail("https://example.com/image.jpg", Detail::Low);
1024        match part {
1025            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(image_part) => {
1026                assert_eq!(image_part.image_url.url, "https://example.com/image.jpg");
1027                assert_eq!(image_part.image_url.detail, Some(Detail::Low));
1028                assert_eq!(image_part.r#type, ImageType::ImageUrl);
1029            }
1030            _ => panic!("Expected image part"),
1031        }
1032    }
1033
1034    #[test]
1035    fn test_image_base64_part() {
1036        let part = image_base64_part("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", "image/png");
1037        match part {
1038            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(image_part) => {
1039                assert!(image_part.image_url.url.starts_with("data:image/png;base64,"));
1040                assert_eq!(image_part.image_url.detail, Some(Detail::Auto));
1041                assert_eq!(image_part.r#type, ImageType::ImageUrl);
1042            }
1043            _ => panic!("Expected image part"),
1044        }
1045    }
1046
1047    #[test]
1048    fn test_image_base64_part_with_detail() {
1049        let part = image_base64_part_with_detail("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", "image/jpeg", Detail::High);
1050        match part {
1051            ChatCompletionRequestUserMessageContentPart::ChatCompletionRequestMessageContentPartImage(image_part) => {
1052                assert!(image_part.image_url.url.starts_with("data:image/jpeg;base64,"));
1053                assert_eq!(image_part.image_url.detail, Some(Detail::High));
1054                assert_eq!(image_part.r#type, ImageType::ImageUrl);
1055            }
1056            _ => panic!("Expected image part"),
1057        }
1058    }
1059
1060    #[test]
1061    fn test_tool_function_with_empty_parameters() {
1062        let tool = tool_function(
1063            "simple_function",
1064            "A simple function",
1065            serde_json::json!({}),
1066        );
1067
1068        assert_eq!(tool.function.name, "simple_function");
1069        assert!(tool.function.parameters.is_some());
1070        assert!(tool.function.parameters.as_ref().unwrap().is_empty());
1071    }
1072
1073    #[test]
1074    fn test_tool_function_with_invalid_parameters() {
1075        let tool = tool_function(
1076            "function_with_string_params",
1077            "A function with string parameters",
1078            serde_json::json!("not an object"),
1079        );
1080
1081        assert_eq!(tool.function.name, "function_with_string_params");
1082        assert!(tool.function.parameters.is_some());
1083        // Should result in empty map when parameters is not an object
1084        assert!(tool.function.parameters.as_ref().unwrap().is_empty());
1085    }
1086}
openai_ergonomic/builders/chat.rs

openai_ergonomic/builders/
chat.rs
