pull private model in embedding svc (#53)

* handle api key for HF

* parse on GET

* pass header to embedding svc

* refactor

Author: ChuckHend

Cargo.toml

@@ -26,12 +26,13 @@ pgrx = "0.11.3"
 postgres-types = "0.2.5"
 regex = "1.9.2"
 reqwest = {version = "0.11.18", features = ["json"] }
-serde = "1.0.173"
+serde = { version = "1.0.173", features = ["derive"] }
 serde_json = "1.0.103"
-sqlx = { version = "0.7.2", features = [
+sqlx = { version = "0.7.3", features = [
     "runtime-tokio-native-tls",
     "postgres",
     "chrono",
+    "json"
 ] }
 thiserror = "1.0.44"
 tiktoken-rs = "0.5.7"

src/api.rs

@@ -1,9 +1,7 @@
 use crate::chat::call_chat;
-use crate::executor::VectorizeMeta;
-use crate::search::{cosine_similarity_search, init_table};
+use crate::search::{self, init_table};
 use crate::transformers::transform;
 use crate::types;
-use crate::util;
 
 use anyhow::Result;
 use pgrx::prelude::*;
@@ -41,41 +39,14 @@ fn table(
 
 #[pg_extern]
 fn search(
-    job_name: &str,
-    query: &str,
+    job_name: String,
+    query: String,
     api_key: default!(Option<String>, "NULL"),
     return_columns: default!(Vec<String>, "ARRAY['*']::text[]"),
     num_results: default!(i32, 10),
 ) -> Result<TableIterator<'static, (name!(search_results, pgrx::JsonB),)>> {
-    let project_meta: VectorizeMeta = if let Ok(Some(js)) = util::get_vectorize_meta_spi(job_name) {
-        js
-    } else {
-        error!("failed to get project metadata");
-    };
-    let proj_params: types::JobParams = serde_json::from_value(
-        serde_json::to_value(project_meta.params).unwrap_or_else(|e| {
-            error!("failed to serialize metadata: {}", e);
-        }),
-    )
-    .unwrap_or_else(|e| error!("failed to deserialize metadata: {}", e));
-
-    let schema = proj_params.schema;
-    let table = proj_params.table;
-
-    let embeddings = transform(query, &project_meta.transformer, api_key);
-
-    let search_results = match project_meta.search_alg {
-        types::SimilarityAlg::pgv_cosine_similarity => cosine_similarity_search(
-            job_name,
-            &schema,
-            &table,
-            &return_columns,
-            num_results,
-            &embeddings[0],
-        )?,
-    };
-
-    Ok(TableIterator::new(search_results))
+    let search_results = search::search(&job_name, &query, api_key, return_columns, num_results)?;
+    Ok(TableIterator::new(search_results.into_iter().map(|r| (r,))))
 }
 
 #[pg_extern]
@@ -127,10 +98,10 @@ fn rag(
     // chat models: currently only supports gpt 3.5 and 4
     // https://platform.openai.com/docs/models/gpt-3-5-turbo
     // https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo
-    chat_model: default!(&str, "'gpt-3.5-turbo'"),
+    chat_model: default!(String, "'gpt-3.5-turbo'"),
     // points to the type of prompt template to use
-    task: default!(&str, "'question_answer'"),
-    api_key: default!(Option<&str>, "NULL"),
+    task: default!(String, "'question_answer'"),
+    api_key: default!(Option<String>, "NULL"),
     // number of records to include in the context
     num_context: default!(i32, 2),
     // truncates context to fit the model's context window
@@ -139,8 +110,8 @@ fn rag(
     let resp = call_chat(
         agent_name,
         query,
-        chat_model,
-        task,
+        &chat_model,
+        &task,
         api_key,
         num_context,
         force_trim,

src/chat.rs

@@ -1,5 +1,6 @@
 use crate::executor::VectorizeMeta;
 use crate::guc;
+use crate::search;
 use crate::types;
 use crate::util::get_vectorize_meta_spi;
 
@@ -39,7 +40,7 @@ pub fn call_chat(
     query: &str,
     chat_model: &str,
     task: &str,
-    api_key: Option<&str>,
+    api_key: Option<String>,
     num_context: i32,
     force_trim: bool,
 ) -> Result<ChatResponse> {
@@ -60,53 +61,28 @@ pub fn call_chat(
     let content_column = job_params.columns[0].clone();
     let pk = job_params.primary_key;
     let columns = vec![pk.clone(), content_column.clone()];
-    // query the relevant vectorize table using the query
-    // TODO: refactor so we can call an internal access vector search function
-    let search_results: Result<Vec<ContextualSearch>, spi::Error> = Spi::connect(|c| {
-        let mut results: Vec<ContextualSearch> = Vec::new();
-        let q = format!(
-            "
-        select search_results from vectorize.search(
-            job_name => '{agent_name}',
-            query => '{query}',
-            return_columns => $1,
-            num_results => {num_context}
-        )",
-        );
-        let tup_table = c.select(
-            &q,
-            None,
-            Some(vec![(
-                PgBuiltInOids::TEXTARRAYOID.oid(),
-                columns.into_datum(),
-            )]),
-        )?;
-
-        for row in tup_table {
-            let row_pgrx_js: pgrx::JsonB = row.get_by_name("search_results").unwrap().unwrap();
-            let row_js: serde_json::Value = row_pgrx_js.0;
-
-            let record_id = row_js
-                .get(&pk)
-                .unwrap_or_else(|| error!("`{pk}` not found"));
-            let content = row_js
-                .get(&content_column)
-                .unwrap_or_else(|| error!("`{content_column}` not found"));
-            let text_content =
-                serde_json::to_string(content).expect("failed to serialize content to string");
-
-            let token_ct = bpe.encode_ordinary(&text_content).len() as i32;
-            results.push(ContextualSearch {
-                record_id: serde_json::to_string(record_id)
-                    .expect("failed to serialize record_id to string"),
-                content: text_content,
-                token_ct,
-            });
-        }
-        Ok(results)
-    });
 
-    let search_results = search_results?;
+    let raw_search = search::search(agent_name, query, api_key.clone(), columns, num_context)?;
+
+    let mut search_results: Vec<ContextualSearch> = Vec::new();
+    for s in raw_search {
+        let row_js: serde_json::Value = s.0;
+        let record_id = row_js
+            .get(&pk)
+            .unwrap_or_else(|| error!("`{pk}` not found"));
+        let content = row_js
+            .get(&content_column)
+            .unwrap_or_else(|| error!("`{content_column}` not found"));
+        let text_content =
+            serde_json::to_string(content).expect("failed to serialize content to string");
+        let token_ct = bpe.encode_ordinary(&text_content).len() as i32;
+        search_results.push(ContextualSearch {
+            record_id: serde_json::to_string(record_id)
+                .expect("failed to serialize record_id to string"),
+            content: text_content,
+            token_ct,
+        });
+    }
 
     // read prompt template
     let res_prompts: Result<PromptTemplate, spi::Error> = Spi::connect(|c| {
@@ -165,7 +141,7 @@ fn render_user_message(user_prompt_template: &str, context: &str, query: &str) -
 fn call_chat_completions(
     prompts: RenderedPrompt,
     model: &str,
-    api_key: Option<&str>,
+    api_key: Option<String>,
 ) -> Result<String> {
     let openai_key = match api_key {
         Some(k) => k.to_string(),

src/init.rs

@@ -70,14 +70,15 @@ pub fn init_embedding_table_query(
     table: &str,
     transformer: &str,
     transform_method: &TableMethod,
+    api_key: Option<String>,
 ) -> Vec<String> {
     check_input(job_name).expect("invalid job name");
     let col_type = match transformer {
         // https://platform.openai.com/docs/guides/embeddings/what-are-embeddings
         // for anything but OpenAI, first call info endpoint to get the embedding dim of the model
         "text-embedding-ada-002" => "vector(1536)".to_owned(),
         _ => {
-            let model_info: TransformerMetadata = sync_get_model_info(transformer)
+            let model_info: TransformerMetadata = sync_get_model_info(transformer, api_key)
                 .expect("failed to call vectorize.embedding_service_url");
             let dim = model_info.embedding_dimension;
             format!("vector({dim})")

src/search.rs

@@ -4,8 +4,10 @@ use crate::init::{self, VECTORIZE_QUEUE};
 use crate::job::{create_insert_trigger, create_trigger_handler, create_update_trigger};
 use crate::transformers::http_handler::sync_get_model_info;
 use crate::transformers::openai;
+use crate::transformers::transform;
 use crate::transformers::types::Inputs;
 use crate::types;
+use crate::util;
 
 use anyhow::Result;
 use pgrx::prelude::*;
@@ -28,15 +30,16 @@ pub fn init_table(
 ) -> Result<String> {
     let job_type = types::JobType::Columns;
 
-    // write job to table
-    let init_job_q = init::init_job_query();
     let arguments = match serde_json::to_value(args) {
         Ok(a) => a,
         Err(e) => {
             error!("invalid json for argument `args`: {}", e);
         }
     };
-    let api_key = arguments.get("api_key");
+    let api_key = match arguments.get("api_key") {
+        Some(k) => Some(serde_json::from_value::<String>(k.clone())?),
+        None => None,
+    };
 
     // get prim key type
     let pkey_type = init::get_column_datatype(schema, table, primary_key);
@@ -46,8 +49,8 @@ pub fn init_table(
     // key can be set in a GUC, so if its required but not provided in args, and not in GUC, error
     match transformer {
         "text-embedding-ada-002" => {
-            let openai_key = match api_key {
-                Some(k) => serde_json::from_value::<String>(k.clone())?,
+            let openai_key = match api_key.clone() {
+                Some(k) => k,
                 None => match guc::get_guc(guc::VectorizeGuc::OpenAIKey) {
                     Some(k) => k,
                     None => {
@@ -59,7 +62,7 @@ pub fn init_table(
         }
         t => {
             // make sure transformer exists
-            let _ = sync_get_model_info(t).expect("transformer does not exist");
+            let _ = sync_get_model_info(t, api_key.clone()).expect("transformer does not exist");
         }
     }
 
@@ -71,12 +74,13 @@ pub fn init_table(
         table_method: table_method.clone(),
         primary_key: primary_key.to_string(),
         pkey_type,
-        api_key: api_key
-            .map(|k| serde_json::from_value::<String>(k.clone()).expect("error parsing api key")),
+        api_key: api_key.clone(),
     };
     let params =
         pgrx::JsonB(serde_json::to_value(valid_params.clone()).expect("error serializing params"));
 
+    // write job to table
+    let init_job_q = init::init_job_query();
     // using SPI here because it is unlikely that this code will be run anywhere but inside the extension.
     // background worker will likely be moved to an external container or service in near future
     let ran: Result<_, spi::Error> = Spi::connect(|mut c| {
@@ -110,8 +114,14 @@ pub fn init_table(
     if ran.is_err() {
         error!("error creating job");
     }
-    let init_embed_q =
-        init::init_embedding_table_query(job_name, schema, table, transformer, &table_method);
+    let init_embed_q = init::init_embedding_table_query(
+        job_name,
+        schema,
+        table,
+        transformer,
+        &table_method,
+        api_key,
+    );
 
     let ran: Result<_, spi::Error> = Spi::connect(|mut c| {
         for q in init_embed_q {
@@ -198,14 +208,50 @@ pub fn init_table(
     Ok(format!("Successfully created job: {job_name}"))
 }
 
+pub fn search(
+    job_name: &str,
+    query: &str,
+    api_key: Option<String>,
+    return_columns: Vec<String>,
+    num_results: i32,
+) -> Result<Vec<pgrx::JsonB>> {
+    let project_meta: VectorizeMeta = if let Ok(Some(js)) = util::get_vectorize_meta_spi(job_name) {
+        js
+    } else {
+        error!("failed to get project metadata");
+    };
+    let proj_params: types::JobParams = serde_json::from_value(
+        serde_json::to_value(project_meta.params).unwrap_or_else(|e| {
+            error!("failed to serialize metadata: {}", e);
+        }),
+    )
+    .unwrap_or_else(|e| error!("failed to deserialize metadata: {}", e));
+
+    let schema = proj_params.schema;
+    let table = proj_params.table;
+
+    let embeddings = transform(query, &project_meta.transformer, api_key);
+
+    match project_meta.search_alg {
+        types::SimilarityAlg::pgv_cosine_similarity => cosine_similarity_search(
+            job_name,
+            &schema,
+            &table,
+            &return_columns,
+            num_results,
+            &embeddings[0],
+        ),
+    }
+}
+
 pub fn cosine_similarity_search(
     project: &str,
     schema: &str,
     table: &str,
     return_columns: &[String],
     num_results: i32,
     embeddings: &[f64],
-) -> Result<Vec<(pgrx::JsonB,)>, spi::Error> {
+) -> Result<Vec<pgrx::JsonB>> {
     let query = format!(
         "
     SELECT to_jsonb(t)
@@ -222,7 +268,8 @@ pub fn cosine_similarity_search(
         cols = return_columns.join(", "),
     );
     Spi::connect(|client| {
-        let mut results: Vec<(pgrx::JsonB,)> = Vec::new();
+        // let mut results: Vec<(pgrx::JsonB,)> = Vec::new();
+        let mut results: Vec<pgrx::JsonB> = Vec::new();
         let tup_table = client.select(
             &query,
             None,
@@ -233,7 +280,7 @@ pub fn cosine_similarity_search(
         )?;
         for row in tup_table {
             match row["results"].value()? {
-                Some(r) => results.push((r,)),
+                Some(r) => results.push(r),
                 None => error!("failed to get results"),
             }
         }