feat: add ai_text_completion function (#10880)

Author: BohuTANG

src/query/functions/src/scalars/vector.rs

@@ -25,6 +25,8 @@ use common_openai::OpenAI;
 use common_vector::cosine_distance;
 
 pub fn register(registry: &mut FunctionRegistry) {
+    // cosine_distance
+    // 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,
@@ -48,6 +50,9 @@ 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>, _, _>(
         "embedding_vector",
         |_, _| FunctionDomain::MayThrow,
@@ -63,11 +68,39 @@ pub fn register(registry: &mut FunctionRegistry) {
                         output.push(result.into());
                     }
                     Err(e) => {
-                        ctx.set_error(output.len(), format!("openai request error:{:?}", e));
+                        ctx.set_error(output.len(), format!("openai embedding request error:{:?}", e));
                         output.push(vec![F32::from(0.0)].into());
                     }
                 }
             },
         ),
     );
+
+    // 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, _, _>(
+        "text_completion",
+        |_, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<StringType, StringType, StringType>(
+            |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());
+                match result {
+                    Ok((resp, _)) => {
+                        output.put_str(&resp);
+                    }
+                    Err(e) => {
+                        ctx.set_error(
+                            output.len(),
+                            format!("openai completion request error:{:?}", e),
+                        );
+                        output.put_str("");
+                    }
+                }
+            },
+        ),
+    );
 }

src/query/functions/tests/it/scalars/testdata/function_list.txt

@@ -2996,6 +2996,8 @@ Functions overloads:
 3 subtract_years(Timestamp NULL, Int64 NULL) :: Timestamp NULL
 0 tan(Float64) :: Float64
 1 tan(Float64 NULL) :: Float64 NULL
+0 text_completion(String, String) :: String
+1 text_completion(String NULL, String NULL) :: String NULL
 0 time_slot(Timestamp) :: Timestamp
 1 time_slot(Timestamp NULL) :: Timestamp NULL
 0 to_base64(String) :: String

src/query/sql/src/planner/semantic/type_check.rs

@@ -1561,6 +1561,7 @@ impl<'a> TypeChecker<'a> {
             "coalesce",
             "last_query_id",
             "ai_embedding_vector",
+            "ai_text_completion",
         ]
     }
 
@@ -1781,6 +1782,28 @@ impl<'a> TypeChecker<'a> {
                         .await,
                 )
             }
+            ("ai_text_completion", args) => {
+                // ai_text_completion(prompt) -> text_completion(prompt, api_key)
+                if args.len() != 1 {
+                    return Some(Err(ErrorCode::BadArguments(
+                        "ai_text_completion(STRING) only accepts one STRING argument",
+                    )
+                    .set_span(span)));
+                }
+
+                // Prompt.
+                let arg1 = args[0];
+                // API key.
+                let arg2 = &Expr::Literal {
+                    span,
+                    lit: Literal::String(GlobalConfig::instance().query.openai_api_key.clone()),
+                };
+
+                Some(
+                    self.resolve_function(span, "text_completion", vec![], &[arg1, arg2])
+                        .await,
+                )
+            }
             _ => None,
         }
     }