chore: change completion model to GPT 3.5 Turbo (#10945)

* chore: change completion model to GPT 3.5 Turbo

* add unit test for openai completion

* add trace to openai api

* fix chat text completion response content

* change the max token from 512 to 1024

Author: BohuTANG

Cargo.lock

@@ -19,6 +19,7 @@ common-exception = { path = "../exception" }
 # GitHub dependencies
 
 # Crates.io dependencies
+log = "0.4"
 metrics = "0.20.1"
 openai_api_rust = { git = "https://github.com/datafuse-extras/openai-api", rev = "5f977a4" }
 

src/common/openai/Cargo.toml

@@ -14,29 +14,18 @@
 
 use common_exception::ErrorCode;
 use common_exception::Result;
+use log::trace;
 use openai_api_rust::completions::CompletionsApi;
 use openai_api_rust::completions::CompletionsBody;
 use openai_api_rust::Auth;
 
 use crate::metrics::metrics_completion_count;
 use crate::metrics::metrics_completion_token;
+use crate::AIModel;
 use crate::OpenAI;
 
-pub enum CompletionMode {
-    // SQL translate:
-    // max_tokens: 150, stop: ['#', ';']
-    SQL,
-    // Text completion:
-    // max_tokens: 512, stop: none
-    Text,
-}
-
 impl OpenAI {
-    pub fn completion_request(
-        &self,
-        prompt: String,
-        mode: CompletionMode,
-    ) -> Result<(String, Option<u32>)> {
+    pub fn completion_sql_request(&self, prompt: String) -> Result<(String, Option<u32>)> {
         let openai = openai_api_rust::OpenAI::new(
             Auth {
                 api_key: self.api_key.clone(),
@@ -45,19 +34,16 @@ impl OpenAI {
             &self.api_base,
         );
 
-        let (max_tokens, stop) = match mode {
-            CompletionMode::SQL => (Some(150), Some(vec!["#".to_string(), ";".to_string()])),
-            CompletionMode::Text => (Some(512), None),
-        };
+        let (max_tokens, stop) = (Some(150), Some(vec!["#".to_string(), ";".to_string()]));
 
         let body = CompletionsBody {
-            model: self.model.to_string(),
+            model: AIModel::TextDavinci003.to_string(),
             prompt: Some(vec![prompt]),
             suffix: None,
             max_tokens,
             temperature: Some(0_f32),
             top_p: Some(1_f32),
-            n: Some(2),
+            n: None,
             stream: Some(false),
             logprobs: None,
             echo: None,
@@ -68,10 +54,14 @@ impl OpenAI {
             logit_bias: None,
             user: None,
         };
+        trace!("openai sql completion request: {:?}", body);
+
         let resp = openai.completion_create(&body).map_err(|e| {
-            ErrorCode::Internal(format!("openai completion request error: {:?}", e))
+            ErrorCode::Internal(format!("openai completion request sql error: {:?}", e))
         })?;
 
+        trace!("openai sql completion response: {:?}", resp);
+
         let usage = resp.usage.total_tokens;
         let sql = if resp.choices.is_empty() {
             "".to_string()

src/common/openai/src/completion.rs src/common/openai/src/completion_sql.rssrc/common/openai/src/completion.rs renamed to src/common/openai/src/completion_sql.rs

@@ -0,0 +1,89 @@
+//  Copyright 2023 Datafuse Labs.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+
+use common_exception::ErrorCode;
+use common_exception::Result;
+use log::trace;
+use openai_api_rust::chat::ChatApi;
+use openai_api_rust::chat::ChatBody;
+use openai_api_rust::Auth;
+use openai_api_rust::Message;
+use openai_api_rust::Role;
+
+use crate::metrics::metrics_completion_count;
+use crate::metrics::metrics_completion_token;
+use crate::AIModel;
+use crate::OpenAI;
+
+impl OpenAI {
+    pub fn completion_text_request(&self, prompt: String) -> Result<(String, Option<u32>)> {
+        let openai = openai_api_rust::OpenAI::new(
+            Auth {
+                api_key: self.api_key.clone(),
+                organization: None,
+            },
+            &self.api_base,
+        );
+
+        let (max_tokens, stop) = (Some(1024), None);
+
+        let body = ChatBody {
+            model: AIModel::GPT35Turbo.to_string(),
+            temperature: Some(0_f32),
+            top_p: Some(1_f32),
+            n: None,
+            stream: None,
+            stop,
+            max_tokens,
+            presence_penalty: None,
+            frequency_penalty: None,
+            logit_bias: None,
+            user: None,
+            messages: vec![Message {
+                role: Role::User,
+                content: prompt,
+            }],
+        };
+
+        trace!("openai text completion request: {:?}", body);
+
+        let resp = openai.chat_completion_create(&body).map_err(|e| {
+            ErrorCode::Internal(format!("openai completion text request error: {:?}", e))
+        })?;
+        trace!("openai text completion response: {:?}", resp);
+
+        let usage = resp.usage.total_tokens;
+        let result = if resp.choices.is_empty() {
+            "".to_string()
+        } else {
+            let message = resp
+                .choices
+                .get(0)
+                .and_then(|choice| choice.message.as_ref());
+
+            match message {
+                Some(msg) => msg.content.clone(),
+                _ => "".to_string(),
+            }
+        };
+
+        // perf.
+        {
+            metrics_completion_count(1);
+            metrics_completion_token(usage.unwrap_or(0));
+        }
+
+        Ok((result, usage))
+    }
+}

src/common/openai/src/completion_text.rs

@@ -20,6 +20,7 @@ use openai_api_rust::Auth;
 
 use crate::metrics::metrics_embedding_count;
 use crate::metrics::metrics_embedding_token;
+use crate::AIModel;
 use crate::OpenAI;
 
 impl OpenAI {
@@ -33,7 +34,7 @@ impl OpenAI {
             &self.api_base,
         );
         let body = EmbeddingsBody {
-            model: self.model.to_string(),
+            model: AIModel::TextEmbeddingAda003.to_string(),
             input: input.to_vec(),
             user: None,
         };

src/common/openai/src/embedding.rs

@@ -12,13 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod completion;
+mod completion_sql;
+mod completion_text;
 mod embedding;
+
 #[allow(clippy::module_inception)]
 mod openai;
 
 pub(crate) mod metrics;
 
-pub use completion::CompletionMode;
 pub use openai::AIModel;
 pub use openai::OpenAI;

src/common/openai/src/lib.rs

@@ -13,8 +13,12 @@
 //  limitations under the License.
 
 pub enum AIModel {
+    // For SQL completion.
     TextDavinci003,
+    // For embedding.
     TextEmbeddingAda003,
+    // For Text completion.
+    GPT35Turbo,
 }
 
 // https://platform.openai.com/examples
@@ -23,22 +27,21 @@ impl ToString for AIModel {
         match self {
             AIModel::TextDavinci003 => "text-davinci-003".to_string(),
             AIModel::TextEmbeddingAda003 => "text-embedding-ada-002".to_string(),
+            AIModel::GPT35Turbo => "gpt-3.5-turbo".to_string(),
         }
     }
 }
 
 pub struct OpenAI {
     pub(crate) api_key: String,
     pub(crate) api_base: String,
-    pub(crate) model: AIModel,
 }
 
 impl OpenAI {
-    pub fn create(api_key: String, model: AIModel) -> Self {
+    pub fn create(api_key: String) -> Self {
         OpenAI {
             api_key,
             api_base: "https://api.openai.com/v1/".to_string(),
-            model,
         }
     }
 }

src/common/openai/src/openai.rs

@@ -0,0 +1,15 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+mod openai;

src/common/openai/tests/it/main.rs

@@ -0,0 +1,55 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_openai::OpenAI;
+
+fn create_openai() -> Option<OpenAI> {
+    let key = std::env::var("OPENAI_API_KEY").unwrap_or("".to_string());
+    if !key.is_empty() {
+        Some(OpenAI::create(key))
+    } else {
+        None
+    }
+}
+
+#[test]
+fn test_openai_text_completion() {
+    let openai = create_openai();
+    if let Some(openai) = openai {
+        let resp = openai
+            .completion_text_request("say hello".to_string())
+            .unwrap();
+
+        assert!(resp.0.contains("hello"));
+    }
+}
+
+#[test]
+fn test_openai_sql_completion() {
+    let openai = create_openai();
+    if let Some(openai) = openai {
+        let resp = openai
+            .completion_sql_request("### Postgres SQL tables, with their properties:
+#
+# Employee(id, name, department_id)
+# Department(id, name, address)
+# Salary_Payments(id, employee_id, amount, date)
+#
+### A query to list the names of the departments which employed more than 10 employees in the last 3 months
+SELECT".to_string())
+            .unwrap();
+
+        assert!(resp.0.contains("FROM"));
+    }
+}

src/common/openai/tests/it/openai.rs

@@ -20,8 +20,6 @@ use common_expression::types::F32;
 use common_expression::vectorize_with_builder_2_arg;
 use common_expression::FunctionDomain;
 use common_expression::FunctionRegistry;
-use common_openai::AIModel;
-use common_openai::CompletionMode;
 use common_openai::OpenAI;
 use common_vector::cosine_distance;
 
@@ -61,7 +59,7 @@ pub fn register(registry: &mut FunctionRegistry) {
             |data, api_key, output, ctx| {
                 let data = std::str::from_utf8(data).unwrap();
                 let api_key = std::str::from_utf8(api_key).unwrap();
-                let openai = OpenAI::create(api_key.to_string(), AIModel::TextEmbeddingAda003);
+                let openai = OpenAI::create(api_key.to_string());
                 let result = openai.embedding_request(&[data.to_string()]);
                 match result {
                     Ok((embeddings, _)) => {
@@ -87,8 +85,8 @@ pub fn register(registry: &mut FunctionRegistry) {
             |data, api_key, output, ctx| {
                 let data = std::str::from_utf8(data).unwrap();
                 let api_key = std::str::from_utf8(api_key).unwrap();
-                let openai = OpenAI::create(api_key.to_string(), AIModel::TextDavinci003);
-                let result = openai.completion_request(data.to_string(), CompletionMode::Text);
+                let openai = OpenAI::create(api_key.to_string());
+                let result = openai.completion_text_request(data.to_string());
                 match result {
                     Ok((resp, _)) => {
                         output.put_str(&resp);