anchor_chain/models/
openai.rs

1//! Module for integrating OpenAI models.
2//!
3//! Facilitates the construction and execution of requests to OpenAI models,
4//! leveraging the OpenAI API.
5
6use std::fmt;
7
8use async_openai::types::{
9    ChatCompletionRequestSystemMessageArgs, ChatCompletionRequestUserMessageArgs,
10    ChatCompletionRequestUserMessageContent, CreateChatCompletionRequestArgs,
11    CreateCompletionRequestArgs, CreateEmbeddingRequestArgs, Prompt,
12};
13use async_trait::async_trait;
14#[cfg(feature = "tracing")]
15use tracing::instrument;
16
17use crate::error::AnchorChainError;
18use crate::models::embedding_model::EmbeddingModel;
19use crate::node::Node;
20
21/// OpenAI model types supported by the `OpenAI` node
22#[derive(Debug)]
23pub enum OpenAIModel<T>
24where
25    T: Into<Prompt>,
26{
27    /// GPT-3.5 Turbo model
28    GPT3_5Turbo(OpenAIChatModel<T>),
29    /// GPT-3.5 Turbo Instruct model
30    GPT3_5TurboInstruct(OpenAIInstructModel<T>),
31    /// GPT-4 Turbo model
32    GPT4Turbo(OpenAIChatModel<T>),
33}
34
35impl<T> OpenAIModel<T>
36where
37    T: Into<Prompt>,
38{
39    /// Constructs a GPT4 Turbo model with the specified system prompt.
40    ///
41    /// The system prompt is passed in as the first message in the conversation
42    /// using `ChatCompletionRequestSystemMessage`.
43    pub async fn new_gpt4_turbo(system_prompt: String) -> Self {
44        OpenAIModel::GPT3_5Turbo(
45            OpenAIChatModel::new(system_prompt, "gpt-4-turbo-preview".to_string()).await,
46        )
47    }
48
49    /// Constructs a GPT3.5 Turbo model with the specified system prompt.
50    ///
51    /// The system prompt is passed in as the first message in the conversation
52    /// using `ChatCompletionRequestSystemMessage`.
53    pub async fn new_gpt3_5_turbo(system_prompt: String) -> Self {
54        OpenAIModel::GPT4Turbo(
55            OpenAIChatModel::new(system_prompt, "gpt-3.5-turbo".to_string()).await,
56        )
57    }
58
59    /// Constructs a GPT3.5 Turbo Instruct model.
60    pub async fn new_gpt3_5_turbo_instruct() -> Self {
61        OpenAIModel::GPT3_5TurboInstruct(
62            OpenAIInstructModel::new("gpt-3.5-turbo-instruct-0914".to_string()).await,
63        )
64    }
65}
66
67#[async_trait]
68impl<T> Node for OpenAIModel<T>
69where
70    T: Send + Sync + fmt::Debug,
71    T: Into<Prompt> + Into<ChatCompletionRequestUserMessageContent>,
72{
73    /// The input that is converted to a `Prompt` for the OpenAI model.
74    type Input = T;
75    /// The output from the OpenAI model.
76    type Output = String;
77
78    /// Sends the prompt to the OpenAI model and processes the response.
79    #[cfg_attr(feature = "tracing", instrument(skip(self)))]
80    async fn process(&self, input: Self::Input) -> Result<Self::Output, AnchorChainError> {
81        match self {
82            OpenAIModel::GPT3_5Turbo(model) => model.process(input).await,
83            OpenAIModel::GPT4Turbo(model) => model.process(input).await,
84            OpenAIModel::GPT3_5TurboInstruct(model) => model.process(input).await,
85        }
86    }
87}
88
89/// Represents a processor for sending and processing requests to the OpenAI API.
90///
91/// `OpenAIChatModel` encapsulates the functionality required to interact with
92/// the OpenAI Chat API, handling both the construction of requests and the
93/// parsing of responses.
94pub struct OpenAIChatModel<T> {
95    system_prompt: String,
96    model: String,
97    client: async_openai::Client<async_openai::config::OpenAIConfig>,
98    _phantom: std::marker::PhantomData<T>,
99}
100
101impl<T> OpenAIChatModel<T> {
102    /// Constructs a new `OpenAI` processor with the default API configuration.
103    ///
104    /// The OpenAIConfig will try to use the API key from the environment
105    /// variable `OPENAI_API_KEY` by default. The system prompt is passed in
106    /// as the first message in the conversation using
107    /// `ChatCompletionRequestSystemMessage`.
108    ///
109    /// Possible Model Types:
110    /// gpt-3.5-turbo-16k
111    /// davinci-002
112    /// gpt-3.5-turbo-1106
113    /// whisper-1
114    /// dall-e-2
115    /// tts-1-hd-1106
116    /// tts-1-hd
117    /// gpt-4-vision-preview
118    /// gpt-3.5-turbo-0125
119    /// gpt-4-turbo-preview
120    /// gpt-3.5-turbo-0301
121    /// gpt-4-1106-preview
122    /// gpt-3.5-turbo
123    /// gpt-4-0613
124    /// gpt-4-1106-vision-preview
125    /// tts-1
126    /// dall-e-3
127    /// babbage-002
128    /// tts-1-1106
129    /// gpt-4
130    /// gpt-4-0125-preview
131    /// gpt-3.5-turbo-0613
132    /// gpt-3.5-turbo-16k-0613
133    async fn new(system_prompt: String, model: String) -> Self {
134        let config = async_openai::config::OpenAIConfig::new();
135        let client = async_openai::Client::with_config(config);
136        OpenAIChatModel {
137            system_prompt,
138            client,
139            model,
140            _phantom: std::marker::PhantomData,
141        }
142    }
143
144    /// Constructs a new `OpenAI` node using the specified API key.
145    ///
146    /// The system prompt is passed in as the first message in the conversation
147    /// using `ChatCompletionRequestSystemMessage`.
148    pub async fn new_with_key(system_prompt: String, model: String, api_key: String) -> Self {
149        let config = async_openai::config::OpenAIConfig::new().with_api_key(api_key);
150        let client = async_openai::Client::with_config(config);
151        OpenAIChatModel {
152            system_prompt,
153            client,
154            model,
155            _phantom: std::marker::PhantomData,
156        }
157    }
158}
159
160#[async_trait]
161impl<T> Node for OpenAIChatModel<T>
162where
163    T: Into<ChatCompletionRequestUserMessageContent> + fmt::Debug + Send + Sync,
164{
165    type Input = T;
166    type Output = String;
167
168    /// Sends the input to the OpenAI API and processes the response.
169    ///
170    /// Constructs a request based on the input and the system prompt, then parses
171    /// the model's response to extract and return final output.
172    #[cfg_attr(feature = "tracing", instrument(skip(self), fields(model = self.model.as_str(), system_prompt = self.system_prompt.as_str())))]
173    async fn process(&self, input: Self::Input) -> Result<Self::Output, AnchorChainError> {
174        let system_prompt = ChatCompletionRequestSystemMessageArgs::default()
175            .content(self.system_prompt.clone())
176            .build()?
177            .into();
178
179        let input = ChatCompletionRequestUserMessageArgs::default()
180            .content(input)
181            .build()?
182            .into();
183
184        let request = CreateChatCompletionRequestArgs::default()
185            .max_tokens(512u16)
186            .model(&self.model)
187            .messages([system_prompt, input])
188            .build()?;
189
190        let response = self.client.chat().create(request).await?;
191        if response.choices.is_empty() {
192            return Err(AnchorChainError::EmptyResponseError);
193        }
194
195        let content = response
196            .choices
197            .first()
198            .ok_or(AnchorChainError::EmptyResponseError)?
199            .message
200            .clone()
201            .content
202            .ok_or(AnchorChainError::EmptyResponseError)?;
203
204        Ok(content)
205    }
206}
207
208impl<T> fmt::Debug for OpenAIChatModel<T> {
209    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
210        f.debug_struct("OpenAI")
211            .field("system_prompt", &self.system_prompt)
212            .finish()
213    }
214}
215
216/// Node for making requests to OpenAI Instruct models.
217pub struct OpenAIInstructModel<T>
218where
219    T: Into<Prompt>,
220{
221    /// The name of the instruct model.
222    model: String,
223    /// The OpenAI API client.
224    client: async_openai::Client<async_openai::config::OpenAIConfig>,
225    _phantom: std::marker::PhantomData<T>,
226}
227
228impl<T> OpenAIInstructModel<T>
229where
230    T: Into<Prompt>,
231{
232    /// Constructs a new `OpenAI` node with the default API configuration.
233    ///
234    /// The model specified must support the instruct API.
235    ///
236    /// Possible Model Types:
237    /// gpt-3.5-turbo-instruct
238    /// gpt-3.5-turbo-instruct-0914
239    #[allow(dead_code)]
240    async fn new(model: String) -> Self {
241        let config = async_openai::config::OpenAIConfig::new();
242        let client = async_openai::Client::with_config(config);
243        OpenAIInstructModel {
244            client,
245            model,
246            _phantom: std::marker::PhantomData,
247        }
248    }
249
250    /// Constructs a new `OpenAI` processor with a specified API key.
251    ///
252    /// The model specified must support the instruct API.
253    #[allow(dead_code)]
254    pub async fn new_with_key(model: String, api_key: String) -> Self {
255        let config = async_openai::config::OpenAIConfig::new().with_api_key(api_key);
256        let client = async_openai::Client::with_config(config);
257        OpenAIInstructModel {
258            client,
259            model,
260            _phantom: std::marker::PhantomData,
261        }
262    }
263}
264
265#[async_trait]
266impl<T> Node for OpenAIInstructModel<T>
267where
268    T: Into<Prompt> + fmt::Debug + Send + Sync,
269{
270    type Input = T;
271    type Output = String;
272
273    /// Sends the input to the OpenAI API and processes the response.
274    ///
275    /// Constructs a request based on the input and the system prompt, then parses
276    /// the model's response to extract and return the processed content.
277    #[cfg_attr(feature = "tracing", instrument(skip(self), fields(model = self.model.as_str())))]
278    async fn process(&self, input: Self::Input) -> Result<Self::Output, AnchorChainError> {
279        let request = CreateCompletionRequestArgs::default()
280            .model(&self.model)
281            .prompt(input)
282            .temperature(0.8)
283            .max_tokens(512u16)
284            .build()?;
285
286        let response = self.client.completions().create(request).await?;
287
288        let content = response
289            .choices
290            .first()
291            .ok_or(AnchorChainError::EmptyResponseError)?
292            .text
293            .clone();
294
295        Ok(content)
296    }
297}
298
299impl<T> fmt::Debug for OpenAIInstructModel<T>
300where
301    T: Into<Prompt>,
302{
303    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
304        f.debug_struct("OpenAI").finish()
305    }
306}
307
308/// Node for making requests to OpenAI embedding models.
309pub struct OpenAIEmbeddingModel {
310    /// The name of the instruct model.
311    model: String,
312    /// The OpenAI API client.
313    client: async_openai::Client<async_openai::config::OpenAIConfig>,
314}
315
316impl Default for OpenAIEmbeddingModel {
317    fn default() -> Self {
318        OpenAIEmbeddingModel {
319            model: "text-embedding-3-large".to_string(),
320            client: async_openai::Client::with_config(async_openai::config::OpenAIConfig::new()),
321        }
322    }
323}
324
325impl OpenAIEmbeddingModel {
326    /// Constructs a new `OpenAI` node with the default API configuration.
327    ///
328    /// The model specified must support the instruct API.
329    ///
330    /// Possible Model Types:
331    /// text-embedding-3-large
332    /// text-embedding-3-small
333    /// text-embedding-ada-002
334    #[allow(dead_code)]
335    async fn new(model: String) -> Self {
336        let config = async_openai::config::OpenAIConfig::new();
337        let client = async_openai::Client::with_config(config);
338        OpenAIEmbeddingModel { client, model }
339    }
340
341    /// Constructs a new `OpenAI` processor with a specified API key.
342    ///
343    /// The model specified must support the embedding API.
344    #[allow(dead_code)]
345    async fn new_with_key(model: String, api_key: String) -> Self {
346        let config = async_openai::config::OpenAIConfig::new().with_api_key(api_key);
347        let client = async_openai::Client::with_config(config);
348        OpenAIEmbeddingModel { client, model }
349    }
350}
351
352#[async_trait]
353impl Node for OpenAIEmbeddingModel {
354    type Input = Vec<String>;
355    type Output = Vec<Vec<f32>>;
356
357    /// Sends the input to the OpenAI API and processes the response.
358    ///
359    /// Constructs a request based on the input and the system prompt, then parses
360    /// the model's response to extract and return the processed content.
361    #[cfg_attr(feature = "tracing", instrument(skip(self), fields(model = self.model.as_str())))]
362    async fn process(&self, input: Self::Input) -> Result<Self::Output, AnchorChainError> {
363        let request = CreateEmbeddingRequestArgs::default()
364            .model(&self.model)
365            .input(input)
366            .build()?;
367
368        let response = self.client.embeddings().create(request).await?;
369
370        Ok(response
371            .data
372            .iter()
373            .map(|data| data.embedding.clone())
374            .collect())
375    }
376}
377
378#[async_trait]
379impl EmbeddingModel for OpenAIEmbeddingModel {
380    #[cfg_attr(feature = "tracing", instrument(skip(self), fields(model = self.model.as_str())))]
381    async fn embed(&self, input: String) -> Result<Vec<f32>, AnchorChainError> {
382        self.process(vec![input])
383            .await?
384            .first()
385            .ok_or(AnchorChainError::EmptyResponseError)
386            .cloned()
387    }
388
389    fn dimensions(&self) -> usize {
390        3072
391    }
392}
393
394impl fmt::Debug for OpenAIEmbeddingModel {
395    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
396        f.debug_struct("OpenAI").finish()
397    }
398}
anchor_chain/models/openai.rs

anchor_chain/models/
openai.rs
