Merge pull request #10645 from BohuTANG/dev-refine-ai

chore(openai): add openai.rs to prepare more AI functions

Author: BohuTANG

src/query/service/src/table_functions/openai/gpt_to_sql.rs

@@ -14,10 +14,7 @@
 
 use std::any::Any;
 use std::sync::Arc;
-use std::time::Duration;
 
-use async_openai::types::CreateCompletionRequestArgs;
-use async_openai::Client;
 use chrono::NaiveDateTime;
 use chrono::TimeZone;
 use chrono::Utc;
@@ -50,6 +47,9 @@ use common_storages_fuse::TableContext;
 use common_storages_view::view_table::VIEW_ENGINE;
 use tracing::info;
 
+use crate::table_functions::openai::AIModel;
+use crate::table_functions::openai::OpenAI;
+
 pub struct GPT2SQLTable {
     prompt: String,
     api_key: String,
@@ -233,37 +233,16 @@ impl AsyncSource for GPT2SQLSource {
         template.push("#".to_string());
         template.push("SELECT".to_string());
 
-        let model = "code-davinci-002";
-        let timeout = Duration::from_secs(30);
         let prompt = template.join("");
-        let api_key = self.api_key.clone();
         info!("openai request prompt: {}", prompt);
 
-        // Client
-        let http_client = reqwest::ClientBuilder::new()
-            .user_agent("databend")
-            .timeout(timeout)
-            .build()
-            .map_err(|e| ErrorCode::Internal(format!("openai http error: {:?}", e)))?;
-        let client = Client::new()
-            .with_api_key(api_key)
-            .with_http_client(http_client);
-
-        // Request
-        let request = CreateCompletionRequestArgs::default()
-            .model(model)
-            .prompt(prompt)
-            .temperature(0.0)
-            .max_tokens(150_u16)
-            .top_p(1.0)
-            .frequency_penalty(0.0)
-            .presence_penalty(0.0)
-            .stop(["#", ";"])
-            .build()
-            .map_err(|e| ErrorCode::Internal(format!("openai request error: {:?}", e)))?;
-
         // Response.
-        let response = client
+        let api_key = self.api_key.clone();
+        let openai = OpenAI::create(api_key, AIModel::CodeDavinci002);
+        let request = openai.completion_request(prompt)?;
+
+        let response = openai
+            .client()?
             .completions()
             .create(request)
             .await

src/query/service/src/table_functions/openai/mod.rs

@@ -13,5 +13,9 @@
 //  limitations under the License.
 
 mod gpt_to_sql;
+#[allow(clippy::module_inception)]
+mod openai;
 
 pub use gpt_to_sql::GPT2SQLTable;
+pub use openai::AIModel;
+pub use openai::OpenAI;

src/query/service/src/table_functions/openai/openai.rs

@@ -0,0 +1,73 @@
+//  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 std::time::Duration;
+
+use async_openai::types::CreateCompletionRequest;
+use async_openai::types::CreateCompletionRequestArgs;
+use async_openai::Client;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+pub enum AIModel {
+    CodeDavinci002,
+}
+
+// https://platform.openai.com/examples
+impl ToString for AIModel {
+    fn to_string(&self) -> String {
+        match self {
+            AIModel::CodeDavinci002 => "code-davinci-002".to_string(),
+        }
+    }
+}
+
+pub struct OpenAI {
+    api_key: String,
+    model: AIModel,
+}
+
+impl OpenAI {
+    pub fn create(api_key: String, model: AIModel) -> Self {
+        OpenAI { api_key, model }
+    }
+
+    pub fn client(&self) -> Result<Client> {
+        let timeout = Duration::from_secs(30);
+        // Client
+        let http_client = reqwest::ClientBuilder::new()
+            .user_agent("databend")
+            .timeout(timeout)
+            .build()
+            .map_err(|e| ErrorCode::Internal(format!("openai http error: {:?}", e)))?;
+
+        Ok(Client::new()
+            .with_api_key(self.api_key.clone())
+            .with_http_client(http_client))
+    }
+
+    pub fn completion_request(&self, prompt: String) -> Result<CreateCompletionRequest> {
+        CreateCompletionRequestArgs::default()
+            .model(self.model.to_string())
+            .prompt(prompt)
+            .temperature(0.0)
+            .max_tokens(150_u16)
+            .top_p(1.0)
+            .frequency_penalty(0.0)
+            .presence_penalty(0.0)
+            .stop(["#", ";"])
+            .build()
+            .map_err(|e| ErrorCode::Internal(format!("openai completion request error: {:?}", e)))
+    }
+}