feat: add openai api_base_url, completion_model and embedding_model to query config (#10993)

* feat: add openai api_base_url, completion_model and embedding_model to query config

* fix the default value

* fix unit test

* fix unit test

Author: BohuTANG

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

@@ -23,11 +23,28 @@ use openai_api_rust::Role;
 
 use crate::metrics::metrics_completion_count;
 use crate::metrics::metrics_completion_token;
-use crate::AIModel;
 use crate::OpenAI;
 
+#[derive(Debug)]
+enum CompletionMode {
+    Sql,
+    Text,
+}
+
 impl OpenAI {
     pub fn completion_text_request(&self, prompt: String) -> Result<(String, Option<u32>)> {
+        self.completion_request(CompletionMode::Text, prompt)
+    }
+
+    pub fn completion_sql_request(&self, prompt: String) -> Result<(String, Option<u32>)> {
+        self.completion_request(CompletionMode::Sql, prompt)
+    }
+
+    fn completion_request(
+        &self,
+        mode: CompletionMode,
+        prompt: String,
+    ) -> Result<(String, Option<u32>)> {
         let openai = openai_api_rust::OpenAI::new(
             Auth {
                 api_key: self.api_key.clone(),
@@ -36,10 +53,13 @@ impl OpenAI {
             &self.api_base,
         );
 
-        let (max_tokens, stop) = (Some(1024), None);
+        let (max_tokens, stop) = match mode {
+            CompletionMode::Sql => (Some(150), Some(vec!["#".to_string(), ";".to_string()])),
+            CompletionMode::Text => (Some(1024), None),
+        };
 
         let body = ChatBody {
-            model: AIModel::GPT35Turbo.to_string(),
+            model: self.completion_model.to_string(),
             temperature: Some(0_f32),
             top_p: Some(1_f32),
             n: None,
@@ -56,12 +76,15 @@ impl OpenAI {
             }],
         };
 
-        trace!("openai text completion request: {:?}", body);
+        trace!("openai {:?} completion request: {:?}", mode, body);
 
         let resp = openai.chat_completion_create(&body).map_err(|e| {
-            ErrorCode::Internal(format!("openai completion text request error: {:?}", e))
+            ErrorCode::Internal(format!(
+                "openai {:?} completion request error: {:?}",
+                mode, e
+            ))
         })?;
-        trace!("openai text completion response: {:?}", resp);
+        trace!("openai {:?} completion response: {:?}", mode, resp);
 
         let usage = resp.usage.total_tokens;
         let result = if resp.choices.is_empty() {

src/common/openai/src/completion_sql.rs

@@ -20,7 +20,6 @@ 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 {
@@ -34,7 +33,7 @@ impl OpenAI {
             &self.api_base,
         );
         let body = EmbeddingsBody {
-            model: AIModel::TextEmbeddingAda003.to_string(),
+            model: self.embedding_model.to_string(),
             input: input.to_vec(),
             user: None,
         };

src/common/openai/src/embedding.rs

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

src/common/openai/src/lib.rs

@@ -12,36 +12,44 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 
-pub enum AIModel {
-    // For SQL completion.
-    TextDavinci003,
-    // For embedding.
-    TextEmbeddingAda003,
-    // For Text completion.
-    GPT35Turbo,
-}
-
-// https://platform.openai.com/examples
-impl ToString for AIModel {
-    fn to_string(&self) -> String {
-        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) embedding_model: String,
+    pub(crate) completion_model: String,
 }
 
 impl OpenAI {
-    pub fn create(api_key: String) -> Self {
+    pub fn create(
+        api_base: String,
+        api_key: String,
+        embedding_model: String,
+        completion_model: String,
+    ) -> Self {
+        // Check and default.
+        let api_base = if api_base.is_empty() {
+            "https://api.openai.com/v1/".to_string()
+        } else {
+            api_base
+        };
+
+        let embedding_model = if embedding_model.is_empty() {
+            "text-embedding-ada-002".to_string()
+        } else {
+            embedding_model
+        };
+
+        let completion_model = if completion_model.is_empty() {
+            "gpt-3.5-turbo".to_string()
+        } else {
+            completion_model
+        };
+
         OpenAI {
+            api_base,
             api_key,
-            api_base: "https://api.openai.com/v1/".to_string(),
+            embedding_model,
+            completion_model,
         }
     }
 }

src/common/openai/src/openai.rs

@@ -17,7 +17,12 @@ 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))
+        Some(OpenAI::create(
+            "".to_string(),
+            key,
+            "".to_string(),
+            "".to_string(),
+        ))
     } else {
         None
     }

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

@@ -1308,9 +1308,20 @@ pub struct QueryConfig {
     #[clap(long)]
     pub disable_system_table_load: bool,
 
+    #[clap(long, default_value = "https://api.openai.com/v1/")]
+    pub openai_api_base_url: String,
+
     // This will not show in system.configs, put it to mask.rs.
     #[clap(long, default_value = "")]
     pub openai_api_key: String,
+
+    /// https://platform.openai.com/docs/models/embeddings
+    #[clap(long, default_value = "text-embedding-ada-002")]
+    pub openai_api_embedding_model: String,
+
+    /// https://platform.openai.com/docs/guides/chat
+    #[clap(long, default_value = "gpt-3.5-turbo")]
+    pub openai_api_completion_model: String,
 }
 
 impl Default for QueryConfig {
@@ -1372,7 +1383,10 @@ impl TryInto<InnerQueryConfig> for QueryConfig {
             internal_enable_sandbox_tenant: self.internal_enable_sandbox_tenant,
             internal_merge_on_read_mutation: self.internal_merge_on_read_mutation,
             disable_system_table_load: self.disable_system_table_load,
+            openai_api_base_url: self.openai_api_base_url,
             openai_api_key: self.openai_api_key,
+            openai_api_completion_model: self.openai_api_completion_model,
+            openai_api_embedding_model: self.openai_api_embedding_model,
         })
     }
 }
@@ -1446,7 +1460,10 @@ impl From<InnerQueryConfig> for QueryConfig {
             table_cache_bloom_index_filter_count: None,
             table_cache_bloom_index_data_bytes: None,
             disable_system_table_load: inner.disable_system_table_load,
+            openai_api_base_url: inner.openai_api_base_url,
             openai_api_key: inner.openai_api_key,
+            openai_api_completion_model: inner.openai_api_completion_model,
+            openai_api_embedding_model: inner.openai_api_embedding_model,
         }
     }
 }

src/query/config/src/config.rs

@@ -174,7 +174,12 @@ pub struct QueryConfig {
     pub internal_merge_on_read_mutation: bool,
     /// Disable some system load(For example system.configs) for cloud security.
     pub disable_system_table_load: bool,
+
+    /// openai
     pub openai_api_key: String,
+    pub openai_api_base_url: String,
+    pub openai_api_embedding_model: String,
+    pub openai_api_completion_model: String,
 }
 
 impl Default for QueryConfig {
@@ -225,8 +230,11 @@ impl Default for QueryConfig {
             internal_enable_sandbox_tenant: false,
             internal_merge_on_read_mutation: false,
             disable_system_table_load: false,
-            openai_api_key: "".to_string(),
             flight_sql_tls_server_key: "".to_string(),
+            openai_api_base_url: "https://api.openai.com/v1/".to_string(),
+            openai_api_key: "".to_string(),
+            openai_api_completion_model: "gpt-3.5-turbo".to_string(),
+            openai_api_embedding_model: "text-embedding-ada-002".to_string(),
         }
     }
 }

src/query/config/src/inner.rs

@@ -18,6 +18,7 @@ use common_expression::types::Float32Type;
 use common_expression::types::StringType;
 use common_expression::types::F32;
 use common_expression::vectorize_with_builder_2_arg;
+use common_expression::vectorize_with_builder_5_arg;
 use common_expression::FunctionDomain;
 use common_expression::FunctionRegistry;
 use common_openai::OpenAI;
@@ -28,7 +29,7 @@ pub fn register(registry: &mut FunctionRegistry) {
     // This function takes two Float32 arrays as input and computes the cosine distance between them.
     registry.register_passthrough_nullable_2_arg::<ArrayType<Float32Type>, ArrayType<Float32Type>, Float32Type, _, _>(
         "cosine_distance",
-        |_,  _| FunctionDomain::MayThrow,
+        |_, _| FunctionDomain::MayThrow,
         vectorize_with_builder_2_arg::<ArrayType<Float32Type>, ArrayType<Float32Type>,  Float32Type>(
             |lhs, rhs, output, ctx| {
                 let l_f32=
@@ -52,14 +53,18 @@ pub fn register(registry: &mut FunctionRegistry) {
     // embedding_vector
     // This function takes two strings as input, sends an API request to OpenAI, and returns the Float32 array of embeddings.
     // The OpenAI API key is pre-configured during the binder phase, so we rewrite this function and set the API key.
-    registry.register_passthrough_nullable_2_arg::<StringType, StringType, ArrayType<Float32Type>, _, _>(
+    registry.register_passthrough_nullable_5_arg::<StringType, StringType, StringType, StringType, StringType, ArrayType<Float32Type>, _, _>(
         "embedding_vector",
-        |_, _| FunctionDomain::MayThrow,
-        vectorize_with_builder_2_arg::<StringType, StringType, ArrayType<Float32Type>>(
-            |data, api_key, output, ctx| {
+        |_, _, _, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_5_arg::<StringType, StringType, StringType, StringType, StringType, ArrayType<Float32Type>>(
+            |data, api_base,api_key, embedding_model, completion_model, 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());
+
+                let api_base = std::str::from_utf8(api_base).unwrap().to_string();
+                let api_key = std::str::from_utf8(api_key).unwrap().to_string();
+                let embedding_model = std::str::from_utf8(embedding_model).unwrap().to_string();
+                let completion_model= std::str::from_utf8(completion_model).unwrap().to_string();
+                let openai = OpenAI::create(api_base, api_key, embedding_model, completion_model);
                 let result = openai.embedding_request(&[data.to_string()]);
                 match result {
                     Ok((embeddings, _)) => {
@@ -78,14 +83,18 @@ pub fn register(registry: &mut FunctionRegistry) {
     // text_completion
     // This function takes two strings as input, sends an API request to OpenAI, and returns the AI-generated completion as a string.
     // The OpenAI API key is pre-configured during the binder phase, so we rewrite this function and set the API key.
-    registry.register_passthrough_nullable_2_arg::<StringType, StringType, StringType, _, _>(
+    registry.register_passthrough_nullable_5_arg::<StringType,StringType,StringType, StringType, StringType, StringType, _, _>(
         "text_completion",
-        |_, _| FunctionDomain::MayThrow,
-        vectorize_with_builder_2_arg::<StringType, StringType, StringType>(
-            |data, api_key, output, ctx| {
+        |_, _, _, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_5_arg::<StringType, StringType, StringType, StringType, StringType, StringType>(
+                |data, api_base,api_key, embedding_model, completion_model, 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());
+
+                let api_base = std::str::from_utf8(api_base).unwrap().to_string();
+                let api_key = std::str::from_utf8(api_key).unwrap().to_string();
+                let embedding_model = std::str::from_utf8(embedding_model).unwrap().to_string();
+                let completion_model= std::str::from_utf8(completion_model).unwrap().to_string();
+                let openai = OpenAI::create(api_base, api_key, embedding_model, completion_model);
                 let result = openai.completion_text_request(data.to_string());
                 match result {
                     Ok((resp, _)) => {