-
Notifications
You must be signed in to change notification settings - Fork 31
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Arbitrary transformer backends (#26)
* add test * pluggable http handler * switch on search * handle worker * add generic adapter * handle a generic webserver * handle response * update migration script
- Loading branch information
Showing
19 changed files
with
490 additions
and
196 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,29 @@ | ||
use anyhow::{Context, Result}; | ||
|
||
use crate::{ | ||
executor::VectorizeMeta, | ||
guc, | ||
transformers::types::{EmbeddingPayload, EmbeddingRequest, Inputs}, | ||
}; | ||
|
||
use super::openai::trim_inputs; | ||
|
||
pub fn prepare_generic_embedding_request( | ||
job_meta: VectorizeMeta, | ||
inputs: &[Inputs], | ||
) -> Result<EmbeddingRequest> { | ||
let text_inputs = trim_inputs(inputs); | ||
let payload = EmbeddingPayload { | ||
input: text_inputs, | ||
model: job_meta.transformer.to_string(), | ||
}; | ||
|
||
let svc_host = guc::get_guc(guc::VectorizeGuc::EmbeddingServiceUrl) | ||
.context("vectorize.embedding_Service_url is not set")?; | ||
|
||
Ok(EmbeddingRequest { | ||
url: svc_host, | ||
payload, | ||
api_key: None, | ||
}) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,61 @@ | ||
use anyhow::Result; | ||
|
||
use crate::transformers::types::{ | ||
EmbeddingPayload, EmbeddingRequest, EmbeddingResponse, Inputs, PairedEmbeddings, | ||
}; | ||
use pgrx::prelude::*; | ||
|
||
pub async fn handle_response<T: for<'de> serde::Deserialize<'de>>( | ||
resp: reqwest::Response, | ||
method: &'static str, | ||
) -> Result<T> { | ||
if !resp.status().is_success() { | ||
let errmsg = format!( | ||
"Failed to call method '{}', received response with status code:{} and body: {}", | ||
method, | ||
resp.status(), | ||
resp.text().await? | ||
); | ||
warning!("pg-vectorize: error handling response: {}", errmsg); | ||
return Err(anyhow::anyhow!(errmsg)); | ||
} | ||
let value = resp.json::<T>().await?; | ||
Ok(value) | ||
} | ||
|
||
// handle an OpenAI compatible embedding transform request | ||
pub async fn openai_embedding_request(request: EmbeddingRequest) -> Result<Vec<Vec<f64>>> { | ||
log!( | ||
"pg-vectorize: openai request size: {}", | ||
request.payload.input.len() | ||
); | ||
let client = reqwest::Client::new(); | ||
let mut req = client | ||
.post(request.url) | ||
.json::<EmbeddingPayload>(&request.payload) | ||
.header("Content-Type", "application/json"); | ||
if let Some(key) = request.api_key { | ||
req = req.header("Authorization", format!("Bearer {}", key)); | ||
} | ||
let resp = req.send().await?; | ||
let embedding_resp = handle_response::<EmbeddingResponse>(resp, "embeddings").await?; | ||
|
||
let embeddings = embedding_resp | ||
.data | ||
.iter() | ||
.map(|d| d.embedding.clone()) | ||
.collect(); | ||
Ok(embeddings) | ||
} | ||
|
||
// merges the vec of inputs with the embedding responses | ||
pub fn merge_input_output(inputs: Vec<Inputs>, values: Vec<Vec<f64>>) -> Vec<PairedEmbeddings> { | ||
inputs | ||
.into_iter() | ||
.zip(values) | ||
.map(|(input, value)| PairedEmbeddings { | ||
primary_key: input.record_id, | ||
embeddings: value, | ||
}) | ||
.collect() | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1 +1,5 @@ | ||
pub mod generic; | ||
pub mod http_handler; | ||
pub mod openai; | ||
pub mod tembo; | ||
pub mod types; |
Oops, something went wrong.