Replace completion with chat completion

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-async-openai = { version = "0.31.1", features = ["_api", "response-types", "responses", "completions"] }
+async-openai = { version = "0.31.1", features = ["_api", "response-types", "responses", "chat-completion"] }
 async-trait = "0.1.89"
 derive_builder = "0.20.2"
 dotenvy = "0.15.7"

README.md

@@ -27,8 +27,9 @@ Here's a simple example with the Weighted Round Robin (WRR) routing mode:
 // Before running the code, make sure to set your OpenAI API key in the environment variable:
 // export OPENAI_API_KEY="your_openai_api_key"
 
+use tokio::runtime::Runtime;
 use arms::client;
-use arms::types::responses;
+use arms::types::chat;
 
 let config = client::Config::builder()
     .provider("openai")
@@ -51,12 +52,15 @@ let config = client::Config::builder()
     .unwrap();
 
 let mut client = client::Client::new(config);
-let request = responses::CreateResponseArgs::default()
-    .input("give me a poem about nature")
+let request = chat::CreateChatCompletionRequestArgs::default()
+    .messages([
+        chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(),
+        chat::ChatCompletionRequestUserMessage::from("Who won the world series in 2020?").into(),
+    ])
     .build()
     .unwrap();
 
-let response = client.create_response(request).await.unwrap();
+let result = Runtime::new().unwrap().block_on(client.create_completion(request));
 ```
 
 ## Contributing

src/client/client.rs

@@ -4,7 +4,7 @@ use crate::client::config::{Config, ModelName};
 use crate::provider::provider;
 use crate::router::router;
 use crate::types::error::OpenAIError;
-use crate::types::{completions, responses};
+use crate::types::{chat, responses};
 
 pub struct Client {
     providers: HashMap<ModelName, Box<dyn provider::Provider>>,
@@ -37,10 +37,11 @@ impl Client {
         provider.create_response(request).await
     }
 
+    // This is chat completion endpoint.
     pub async fn create_completion(
         &mut self,
-        request: completions::CreateCompletionRequest,
-    ) -> Result<completions::CreateCompletionResponse, OpenAIError> {
+        request: chat::CreateChatCompletionRequest,
+    ) -> Result<chat::CreateChatCompletionResponse, OpenAIError> {
         let candidate = self.router.sample();
         let provider = self.providers.get(&candidate).unwrap();
         provider.create_completion(request).await

src/lib.rs

@@ -14,7 +14,7 @@ mod provider {
     pub mod provider;
 }
 pub mod types {
-    pub mod completions;
+    pub mod chat;
     pub mod error;
     pub mod responses;
 }

src/main.rs

@@ -1,7 +1,7 @@
 use tokio::runtime::Runtime;
 
 use arms::client;
-use arms::types::{completions, responses};
+use arms::types::chat;
 
 fn main() {
     // case 1: completion with DeepInfra provider.
@@ -27,8 +27,17 @@ fn main() {
 
     let mut client = client::Client::new(config);
 
-    let request = completions::CreateCompletionRequestArgs::default()
-        .prompt("How to achieve AGI?")
+    let request = chat::CreateChatCompletionRequestArgs::default()
+        .messages([
+            chat::ChatCompletionRequestSystemMessage::from("You are a helpful assistant.").into(),
+            chat::ChatCompletionRequestUserMessage::from("Who won the world series in 2020?")
+                .into(),
+            // chat::ChatCompletionRequestAssistantMessage::from(
+            //     "The Los Angeles Dodgers won the World Series in 2020.",
+            // )
+            // .into(),
+            // chat::ChatCompletionRequestUserMessage::from("Where was it played?").into(),
+        ])
         .build()
         .unwrap();
 
@@ -38,7 +47,7 @@ fn main() {
 
     match result {
         Ok(response) => {
-            println!("Response ID: {}", response.id);
+            println!("Response: {:?}", response);
         }
         Err(e) => {
             eprintln!("Error: {}", e);

src/provider/common.rs

@@ -1,8 +1,8 @@
 use crate::types::error::OpenAIError;
-use crate::types::{completions, responses};
+use crate::types::{chat, responses};
 
 pub fn validate_completion_request(
-    request: &completions::CreateCompletionRequest,
+    request: &chat::CreateChatCompletionRequest,
 ) -> Result<(), OpenAIError> {
     if request.model != "" {
         return Err(OpenAIError::InvalidArgument(

src/provider/faker.rs

@@ -1,9 +1,8 @@
-use async_openai::types::chat::Choice;
 use async_trait::async_trait;
 
 use crate::client::config::{ModelConfig, ModelName};
 use crate::provider::{common, provider};
-use crate::types::completions::{CreateCompletionRequest, CreateCompletionResponse};
+use crate::types::chat;
 use crate::types::error::OpenAIError;
 use crate::types::responses::{
     AssistantRole, CreateResponse, OutputItem, OutputMessage, OutputMessageContent, OutputStatus,
@@ -76,22 +75,30 @@ impl provider::Provider for FakerProvider {
 
     async fn create_completion(
         &self,
-        _request: CreateCompletionRequest,
-    ) -> Result<CreateCompletionResponse, OpenAIError> {
-        common::validate_completion_request(&_request)?;
-
-        Ok(CreateCompletionResponse {
+        request: chat::CreateChatCompletionRequest,
+    ) -> Result<chat::CreateChatCompletionResponse, OpenAIError> {
+        common::validate_completion_request(&request)?;
+        Ok(chat::CreateChatCompletionResponse {
             id: "fake-completion-id".to_string(),
             object: "text_completion".to_string(),
             created: 1_600_000_000,
             model: self.model.clone(),
-            choices: vec![Choice {
+            usage: None,
+            service_tier: None,
+            choices: vec![chat::ChatChoice {
                 index: 0,
-                text: "This is a fake completion.".to_string(),
-                logprobs: None,
+                message: chat::ChatCompletionResponseMessage {
+                    role: chat::Role::Assistant,
+                    content: Some("This is a fake chat completion.".to_string()),
+                    refusal: None,
+                    tool_calls: None,
+                    annotations: None,
+                    function_call: None,
+                    audio: None,
+                },
                 finish_reason: None,
+                logprobs: None,
             }],
-            usage: None,
             system_fingerprint: None,
         })
     }

src/provider/openai.rs

@@ -5,7 +5,7 @@ use derive_builder::Builder;
 use crate::client::config::{DEFAULT_PROVIDER, ModelConfig, ModelName};
 use crate::provider::{common, provider};
 use crate::types::error::OpenAIError;
-use crate::types::{completions, responses};
+use crate::types::{chat, responses};
 
 #[derive(Debug, Clone, Builder)]
 #[builder(pattern = "mutable", build_fn(skip))]
@@ -65,14 +65,14 @@ impl provider::Provider for OpenAIProvider {
 
     async fn create_completion(
         &self,
-        request: completions::CreateCompletionRequest,
-    ) -> Result<completions::CreateCompletionResponse, OpenAIError> {
+        request: chat::CreateChatCompletionRequest,
+    ) -> Result<chat::CreateChatCompletionResponse, OpenAIError> {
         common::validate_completion_request(&request)?;
 
         // Set the model after validation since model is bind to the provider.
         let mut req = request.clone();
         req.model = self.model.clone();
-        self.client.completions().create(req).await
+        self.client.chat().create(req).await
     }
 
     async fn create_response(

src/provider/provider.rs

@@ -4,7 +4,7 @@ use crate::client::config::ModelConfig;
 use crate::provider::faker::FakerProvider;
 use crate::provider::openai::OpenAIProvider;
 use crate::types::error::OpenAIError;
-use crate::types::{completions, responses};
+use crate::types::{chat, responses};
 
 // Not all providers support response endpoint.
 pub const RESPONSE_ENDPOINT_PROVIDERS: &[&str] = &["FAKER", "OPENAI"];
@@ -33,8 +33,8 @@ pub trait Provider: Send + Sync {
     ) -> Result<responses::Response, OpenAIError>;
     async fn create_completion(
         &self,
-        request: completions::CreateCompletionRequest,
-    ) -> Result<completions::CreateCompletionResponse, OpenAIError>;
+        request: chat::CreateChatCompletionRequest,
+    ) -> Result<chat::CreateChatCompletionResponse, OpenAIError>;
 }
 
 #[cfg(test)]

src/router/router.rs

@@ -1,7 +1,6 @@
 use crate::client::config::{ModelConfig, ModelName, RoutingMode};
 use crate::router::random::RandomRouter;
 use crate::router::wrr::WeightedRoundRobinRouter;
-use crate::types::responses::CreateResponse;
 
 #[derive(Debug, Clone)]
 pub struct ModelInfo {