langchain_rust/llm/openai/
mod.rs

1use std::pin::Pin;
2
3pub use async_openai::config::{AzureConfig, Config, OpenAIConfig};
4
5use async_openai::types::chat::{
6    ChatChoiceStream, ChatCompletionMessageToolCalls, ChatCompletionRequestAssistantMessageArgs,
7    ChatCompletionRequestMessage, ChatCompletionRequestMessageContentPartImageArgs,
8    ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestToolMessageArgs,
9    ChatCompletionRequestUserMessageArgs, ChatCompletionRequestUserMessageContent,
10    ChatCompletionRequestUserMessageContentPart, ChatCompletionStreamOptions,
11    ChatCompletionToolChoiceOption, ChatCompletionTools, CreateChatCompletionRequest,
12    CreateChatCompletionRequestArgs, ResponseFormat, StopConfiguration,
13};
14use async_openai::{error::OpenAIError, Client};
15use async_trait::async_trait;
16use futures::{Stream, StreamExt};
17
18use crate::schemas::convert::{LangchainIntoOpenAI, TryLangchainIntoOpenAI};
19use crate::{
20    language_models::{llm::LLM, options::CallOptions, GenerateResult, LLMError, TokenUsage},
21    schemas::{
22        messages::{Message, MessageType},
23        StreamData, ToolCall,
24    },
25};
26
27#[derive(Clone)]
28pub enum OpenAIModel {
29    Gpt35,
30    Gpt4,
31    Gpt4Turbo,
32    Gpt4o,
33    Gpt4oMini,
34}
35
36impl ToString for OpenAIModel {
37    fn to_string(&self) -> String {
38        match self {
39            OpenAIModel::Gpt35 => "gpt-3.5-turbo".to_string(),
40            OpenAIModel::Gpt4 => "gpt-4".to_string(),
41            OpenAIModel::Gpt4Turbo => "gpt-4-turbo-preview".to_string(),
42            OpenAIModel::Gpt4o => "gpt-4o".to_string(),
43            OpenAIModel::Gpt4oMini => "gpt-4o-mini".to_string(),
44        }
45    }
46}
47
48impl Into<String> for OpenAIModel {
49    fn into(self) -> String {
50        self.to_string()
51    }
52}
53
54#[derive(Clone)]
55pub struct OpenAI<C: Config> {
56    config: C,
57    options: CallOptions,
58    model: String,
59}
60
61impl<C: Config> OpenAI<C> {
62    pub fn new(config: C) -> Self {
63        Self {
64            config,
65            options: CallOptions::default(),
66            model: OpenAIModel::Gpt4oMini.to_string(),
67        }
68    }
69
70    pub fn with_model<S: Into<String>>(mut self, model: S) -> Self {
71        self.model = model.into();
72        self
73    }
74
75    pub fn with_config(mut self, config: C) -> Self {
76        self.config = config;
77        self
78    }
79
80    pub fn with_options(mut self, options: CallOptions) -> Self {
81        self.options = options;
82        self
83    }
84}
85
86impl Default for OpenAI<OpenAIConfig> {
87    fn default() -> Self {
88        Self::new(OpenAIConfig::default())
89    }
90}
91
92#[async_trait]
93impl<C: Config + Clone + Send + Sync + 'static> LLM for OpenAI<C> {
94    async fn generate(&self, prompt: &[Message]) -> Result<GenerateResult, LLMError> {
95        let client = Client::with_config(self.config.clone());
96        let request = self.generate_request(prompt, self.options.streaming_func.is_some())?;
97        match &self.options.streaming_func {
98            Some(func) => {
99                let mut stream = client.chat().create_stream(request).await?;
100                let mut generate_result = GenerateResult::default();
101                while let Some(result) = stream.next().await {
102                    match result {
103                        Ok(response) => {
104                            if let Some(usage) = response.usage {
105                                generate_result.tokens = Some(TokenUsage {
106                                    prompt_tokens: usage.prompt_tokens,
107                                    completion_tokens: usage.completion_tokens,
108                                    total_tokens: usage.total_tokens,
109                                });
110                            }
111                            for chat_choice in response.choices.iter() {
112                                let chat_choice: ChatChoiceStream = chat_choice.clone();
113                                {
114                                    let mut func = func.lock().await;
115                                    let _ = func(
116                                        serde_json::to_string(&chat_choice).unwrap_or("".into()),
117                                    )
118                                    .await;
119                                }
120                                if let Some(content) = chat_choice.delta.content {
121                                    generate_result.generation.push_str(&content);
122                                }
123                            }
124                        }
125                        Err(err) => {
126                            eprintln!("Error from streaming response: {:?}", err);
127                        }
128                    }
129                }
130                Ok(generate_result)
131            }
132            None => {
133                let response = client.chat().create(request).await?;
134                let mut generate_result = GenerateResult::default();
135
136                if let Some(usage) = response.usage {
137                    generate_result.tokens = Some(TokenUsage {
138                        prompt_tokens: usage.prompt_tokens,
139                        completion_tokens: usage.completion_tokens,
140                        total_tokens: usage.total_tokens,
141                    });
142                }
143
144                if let Some(choice) = &response.choices.first() {
145                    generate_result.generation = choice.message.content.clone().unwrap_or_default();
146                    if let Some(calls) = &choice.message.tool_calls {
147                        let typed: Vec<ToolCall> = calls
148                            .iter()
149                            .filter_map(|c| match c {
150                                ChatCompletionMessageToolCalls::Function(f) => Some(ToolCall::new(
151                                    &f.id,
152                                    &f.function.name,
153                                    &f.function.arguments,
154                                )),
155                                _ => None,
156                            })
157                            .collect();
158                        if typed.is_empty() {
159                            // Only non-function (custom) tool calls: keep the raw JSON.
160                            generate_result.generation =
161                                serde_json::to_string(&calls).unwrap_or_default();
162                        } else {
163                            generate_result = generate_result.with_tool_calls(typed);
164                        }
165                    }
166                } else {
167                    generate_result.generation = "".to_string();
168                }
169
170                Ok(generate_result)
171            }
172        }
173    }
174
175    async fn invoke(&self, prompt: &str) -> Result<String, LLMError> {
176        self.generate(&[Message::new_human_message(prompt)])
177            .await
178            .map(|res| res.generation)
179    }
180
181    async fn stream(
182        &self,
183        messages: &[Message],
184    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, LLMError>> + Send>>, LLMError> {
185        let client = Client::with_config(self.config.clone());
186        let request = self.generate_request(messages, true)?;
187
188        let original_stream = client.chat().create_stream(request).await?;
189
190        let new_stream = original_stream.map(|result| match result {
191            Ok(completion) => {
192                let value_completion = serde_json::to_value(completion).map_err(LLMError::from)?;
193                let tokens = match value_completion.pointer("/usage") {
194                    Some(usage) if !usage.is_null() => Some(
195                        serde_json::from_value::<TokenUsage>(usage.clone())
196                            .map_err(LLMError::from)?,
197                    ),
198                    _ => None,
199                };
200                // Some OpenAI-compatible providers send usage together with the last
201                // content delta (#290), so the content is read even when usage is present.
202                let content = value_completion.pointer("/choices/0/delta/content");
203                if tokens.is_none() && content.is_none() {
204                    return Err(LLMError::ContentNotFound(
205                        "/choices/0/delta/content".to_string(),
206                    ));
207                }
208                let content = content
209                    .and_then(serde_json::Value::as_str)
210                    .unwrap_or("")
211                    .to_string();
212
213                Ok(StreamData::new(value_completion, tokens, content))
214            }
215            Err(e) => Err(LLMError::from(e)),
216        });
217
218        Ok(Box::pin(new_stream))
219    }
220
221    fn add_options(&mut self, options: CallOptions) {
222        self.options.merge_options(options)
223    }
224}
225
226impl<C: Config> OpenAI<C> {
227    fn to_openai_messages(
228        &self,
229        messages: &[Message],
230    ) -> Result<Vec<ChatCompletionRequestMessage>, LLMError> {
231        let mut openai_messages: Vec<ChatCompletionRequestMessage> = Vec::new();
232        for m in messages {
233            match m.message_type {
234                MessageType::AIMessage => openai_messages.push(match &m.tool_calls {
235                    Some(value) => {
236                        let function: Vec<ChatCompletionMessageToolCalls> =
237                            serde_json::from_value(value.clone())?;
238                        ChatCompletionRequestAssistantMessageArgs::default()
239                            .tool_calls(function)
240                            .content(m.content.clone())
241                            .build()?
242                            .into()
243                    }
244                    None => ChatCompletionRequestAssistantMessageArgs::default()
245                        .content(m.content.clone())
246                        .build()?
247                        .into(),
248                }),
249                MessageType::HumanMessage => {
250                    let content: ChatCompletionRequestUserMessageContent = match m.images.clone() {
251                        Some(images) => {
252                            let content: Result<
253                                Vec<ChatCompletionRequestUserMessageContentPart>,
254                                OpenAIError,
255                            > = images
256                                .into_iter()
257                                .map(|image| {
258                                    Ok(ChatCompletionRequestMessageContentPartImageArgs::default()
259                                        .image_url(image.image_url)
260                                        .build()?
261                                        .into())
262                                })
263                                .collect();
264
265                            content?.into()
266                        }
267                        None => m.content.clone().into(),
268                    };
269
270                    openai_messages.push(
271                        ChatCompletionRequestUserMessageArgs::default()
272                            .content(content)
273                            .build()?
274                            .into(),
275                    )
276                }
277                MessageType::SystemMessage => openai_messages.push(
278                    ChatCompletionRequestSystemMessageArgs::default()
279                        .content(m.content.clone())
280                        .build()?
281                        .into(),
282                ),
283                MessageType::ToolMessage => {
284                    openai_messages.push(
285                        ChatCompletionRequestToolMessageArgs::default()
286                            .content(m.content.clone())
287                            .tool_call_id(m.id.clone().unwrap_or_default())
288                            .build()?
289                            .into(),
290                    );
291                }
292            }
293        }
294        Ok(openai_messages)
295    }
296
297    fn generate_request(
298        &self,
299        messages: &[Message],
300        stream: bool,
301    ) -> Result<CreateChatCompletionRequest, LLMError> {
302        let messages: Vec<ChatCompletionRequestMessage> = self.to_openai_messages(messages)?;
303        let mut request_builder = CreateChatCompletionRequestArgs::default();
304        if let Some(temperature) = self.options.temperature {
305            request_builder.temperature(temperature);
306        }
307        if let Some(max_tokens) = self.options.max_tokens {
308            request_builder.max_tokens(max_tokens);
309        }
310        if stream {
311            request_builder.stream(true);
312            if let Some(include_usage) = self.options.stream_usage {
313                request_builder.stream_options(ChatCompletionStreamOptions {
314                    include_usage: Some(include_usage),
315                    include_obfuscation: None,
316                });
317            }
318        }
319        request_builder.model(self.model.to_string());
320        if let Some(stop_words) = &self.options.stop_words {
321            request_builder.stop(StopConfiguration::StringArray(stop_words.clone()));
322        }
323
324        if let Some(functions) = &self.options.functions {
325            let functions: Result<Vec<ChatCompletionTools>, OpenAIError> = functions
326                .clone()
327                .into_iter()
328                .map(|f| f.try_into_openai())
329                .collect();
330            request_builder.tools(functions?);
331        }
332
333        if let Some(behavior) = &self.options.function_call_behavior {
334            request_builder
335                .tool_choice::<ChatCompletionToolChoiceOption>(behavior.clone().into_openai());
336        }
337
338        if let Some(response_format) = &self.options.response_format {
339            request_builder
340                .response_format::<ResponseFormat>(response_format.clone().into_openai());
341        }
342
343        request_builder.messages(messages);
344        Ok(request_builder.build()?)
345    }
346}
347
348#[cfg(test)]
349mod mock_tests;
350
351#[cfg(test)]
352mod tests {
353    use crate::schemas::FunctionDefinition;
354
355    use super::*;
356
357    use base64::prelude::*;
358    use serde_json::json;
359    use std::sync::Arc;
360    use tokio::sync::Mutex;
361    use tokio::test;
362
363    #[test]
364    #[ignore]
365    async fn test_invoke() {
366        let message_complete = Arc::new(Mutex::new(String::new()));
367
368        // Define the streaming function
369        // This function will append the content received from the stream to `message_complete`
370        let streaming_func = {
371            let message_complete = message_complete.clone();
372            move |content: String| {
373                let message_complete = message_complete.clone();
374                async move {
375                    let mut message_complete_lock = message_complete.lock().await;
376                    println!("Content: {:?}", content);
377                    message_complete_lock.push_str(&content);
378                    Ok(())
379                }
380            }
381        };
382        let options = CallOptions::new().with_streaming_func(streaming_func);
383        // Setup the OpenAI client with the necessary options
384        let open_ai = OpenAI::new(OpenAIConfig::default())
385            .with_model(OpenAIModel::Gpt35.to_string()) // You can change the model as needed
386            .with_options(options);
387
388        // Define a set of messages to send to the generate function
389
390        // Call the generate function
391        match open_ai.invoke("hola").await {
392            Ok(result) => {
393                // Print the response from the generate function
394                println!("Generate Result: {:?}", result);
395                println!("Message Complete: {:?}", message_complete.lock().await);
396            }
397            Err(e) => {
398                // Handle any errors
399                eprintln!("Error calling generate: {:?}", e);
400            }
401        }
402    }
403
404    #[test]
405    #[ignore]
406    async fn test_generate_function() {
407        let message_complete = Arc::new(Mutex::new(String::new()));
408
409        // Define the streaming function
410        // This function will append the content received from the stream to `message_complete`
411        let streaming_func = {
412            let message_complete = message_complete.clone();
413            move |content: String| {
414                let message_complete = message_complete.clone();
415                async move {
416                    let content = serde_json::from_str::<ChatChoiceStream>(&content).unwrap();
417                    if content.finish_reason.is_some() {
418                        return Ok(());
419                    }
420                    let mut message_complete_lock = message_complete.lock().await;
421                    println!("Content: {:?}", content);
422                    message_complete_lock.push_str(&content.delta.content.unwrap());
423                    Ok(())
424                }
425            }
426        };
427        // Define the streaming function as an async block without capturing external references directly
428        let options = CallOptions::new().with_streaming_func(streaming_func);
429        // Setup the OpenAI client with the necessary options
430        let open_ai = OpenAI::new(OpenAIConfig::default())
431            .with_model(OpenAIModel::Gpt35.to_string()) // You can change the model as needed
432            .with_options(options);
433
434        // Define a set of messages to send to the generate function
435        let messages = vec![Message::new_human_message("Hello, how are you?")];
436
437        // Call the generate function
438        match open_ai.generate(&messages).await {
439            Ok(result) => {
440                // Print the response from the generate function
441                println!("Generate Result: {:?}", result);
442                println!("Message Complete: {:?}", message_complete.lock().await);
443            }
444            Err(e) => {
445                // Handle any errors
446                eprintln!("Error calling generate: {:?}", e);
447            }
448        }
449    }
450
451    #[test]
452    #[ignore]
453    async fn test_openai_stream() {
454        // Setup the OpenAI client with the necessary options
455        let open_ai = OpenAI::default().with_model(OpenAIModel::Gpt35.to_string());
456
457        // Define a set of messages to send to the generate function
458        let messages = vec![Message::new_human_message("Hello, how are you?")];
459
460        open_ai
461            .stream(&messages)
462            .await
463            .unwrap()
464            .for_each(|result| async {
465                match result {
466                    Ok(stream_data) => {
467                        println!("Stream Data: {:?}", stream_data.content);
468                    }
469                    Err(e) => {
470                        eprintln!("Error calling generate: {:?}", e);
471                    }
472                }
473            })
474            .await;
475    }
476
477    #[test]
478    #[ignore]
479    async fn test_function() {
480        let mut functions = Vec::new();
481        functions.push(FunctionDefinition {
482            name: "cli".to_string(),
483            description: "Use the Ubuntu command line to preform any action you wish.".to_string(),
484            parameters: json!({
485                "type": "object",
486                "properties": {
487                    "command": {
488                        "type": "string",
489                        "description": "The raw command you want executed"
490                    }
491                },
492                "required": ["command"]
493            }),
494        });
495
496        let llm = OpenAI::default()
497            .with_model(OpenAIModel::Gpt35)
498            .with_config(OpenAIConfig::new())
499            .with_options(CallOptions::new().with_functions(functions));
500        let response = llm
501            .invoke("Use the command line to create a new rust project. Execute the first command.")
502            .await
503            .unwrap();
504        println!("{}", response)
505    }
506
507    #[test]
508    #[ignore]
509    async fn test_generate_with_image_message() {
510        // Setup the OpenAI client with the necessary options
511        let open_ai =
512            OpenAI::new(OpenAIConfig::default()).with_model(OpenAIModel::Gpt4o.to_string());
513
514        // Convert image to base64
515        let image = std::fs::read("./src/llm/test_data/example.jpg").unwrap();
516        let image_base64 = BASE64_STANDARD.encode(image);
517
518        // Define a set of messages to send to the generate function
519        let image_urls = vec![format!("data:image/jpeg;base64,{image_base64}")];
520        let messages = vec![
521            Message::new_human_message("Describe this image"),
522            Message::new_human_message_with_images(image_urls),
523        ];
524
525        // Call the generate function
526        let response = open_ai.generate(&messages).await.unwrap();
527        println!("Response: {:?}", response);
528    }
529}
langchain_rust/llm/openai/mod.rs

langchain_rust/llm/openai/
mod.rs
