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

openai_ergonomic/builders/
chat.rs
