432 lines
13 KiB
Rust
432 lines
13 KiB
Rust
use anyhow::{anyhow, Result};
|
|
use async_trait::async_trait;
|
|
use futures::AsyncReadExt;
|
|
use gpui::executor::Background;
|
|
use gpui::{serde_json, ViewContext};
|
|
use isahc::http::StatusCode;
|
|
use isahc::prelude::Configurable;
|
|
use isahc::{AsyncBody, Response};
|
|
use lazy_static::lazy_static;
|
|
use ordered_float::OrderedFloat;
|
|
use parking_lot::Mutex;
|
|
use parse_duration::parse;
|
|
use postage::watch;
|
|
use rusqlite::types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef};
|
|
use rusqlite::ToSql;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::env;
|
|
use std::ops::Add;
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tiktoken_rs::{cl100k_base, CoreBPE};
|
|
use util::http::{HttpClient, Request};
|
|
use util::ResultExt;
|
|
|
|
use crate::completion::OPENAI_API_URL;
|
|
|
|
lazy_static! {
|
|
static ref OPENAI_BPE_TOKENIZER: CoreBPE = cl100k_base().unwrap();
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Clone)]
|
|
pub struct Embedding(pub Vec<f32>);
|
|
|
|
// This is needed for semantic index functionality
|
|
// Unfortunately it has to live wherever the "Embedding" struct is created.
|
|
// Keeping this in here though, introduces a 'rusqlite' dependency into AI
|
|
// which is less than ideal
|
|
impl FromSql for Embedding {
|
|
fn column_result(value: ValueRef) -> FromSqlResult<Self> {
|
|
let bytes = value.as_blob()?;
|
|
let embedding: Result<Vec<f32>, Box<bincode::ErrorKind>> = bincode::deserialize(bytes);
|
|
if embedding.is_err() {
|
|
return Err(rusqlite::types::FromSqlError::Other(embedding.unwrap_err()));
|
|
}
|
|
Ok(Embedding(embedding.unwrap()))
|
|
}
|
|
}
|
|
|
|
impl ToSql for Embedding {
|
|
fn to_sql(&self) -> rusqlite::Result<ToSqlOutput> {
|
|
let bytes = bincode::serialize(&self.0)
|
|
.map_err(|err| rusqlite::Error::ToSqlConversionFailure(Box::new(err)))?;
|
|
Ok(ToSqlOutput::Owned(rusqlite::types::Value::Blob(bytes)))
|
|
}
|
|
}
|
|
impl From<Vec<f32>> for Embedding {
|
|
fn from(value: Vec<f32>) -> Self {
|
|
Embedding(value)
|
|
}
|
|
}
|
|
|
|
impl Embedding {
|
|
pub fn similarity(&self, other: &Self) -> OrderedFloat<f32> {
|
|
let len = self.0.len();
|
|
assert_eq!(len, other.0.len());
|
|
|
|
let mut result = 0.0;
|
|
unsafe {
|
|
matrixmultiply::sgemm(
|
|
1,
|
|
len,
|
|
1,
|
|
1.0,
|
|
self.0.as_ptr(),
|
|
len as isize,
|
|
1,
|
|
other.0.as_ptr(),
|
|
1,
|
|
len as isize,
|
|
0.0,
|
|
&mut result as *mut f32,
|
|
1,
|
|
1,
|
|
);
|
|
}
|
|
OrderedFloat(result)
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct OpenAIEmbeddings {
|
|
pub api_key: Option<String>,
|
|
pub client: Arc<dyn HttpClient>,
|
|
pub executor: Arc<Background>,
|
|
rate_limit_count_rx: watch::Receiver<Option<Instant>>,
|
|
rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct OpenAIEmbeddingRequest<'a> {
|
|
model: &'static str,
|
|
input: Vec<&'a str>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct OpenAIEmbeddingResponse {
|
|
data: Vec<OpenAIEmbedding>,
|
|
usage: OpenAIEmbeddingUsage,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct OpenAIEmbedding {
|
|
embedding: Vec<f32>,
|
|
index: usize,
|
|
object: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct OpenAIEmbeddingUsage {
|
|
prompt_tokens: usize,
|
|
total_tokens: usize,
|
|
}
|
|
|
|
#[async_trait]
|
|
pub trait EmbeddingProvider: Sync + Send {
|
|
fn is_authenticated(&self) -> bool;
|
|
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
|
|
fn max_tokens_per_batch(&self) -> usize;
|
|
fn truncate(&self, span: &str) -> (String, usize);
|
|
fn rate_limit_expiration(&self) -> Option<Instant>;
|
|
}
|
|
|
|
pub struct DummyEmbeddings {}
|
|
|
|
#[async_trait]
|
|
impl EmbeddingProvider for DummyEmbeddings {
|
|
fn is_authenticated(&self) -> bool {
|
|
true
|
|
}
|
|
fn rate_limit_expiration(&self) -> Option<Instant> {
|
|
None
|
|
}
|
|
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
|
|
// 1024 is the OpenAI Embeddings size for ada models.
|
|
// the model we will likely be starting with.
|
|
let dummy_vec = Embedding::from(vec![0.32 as f32; 1536]);
|
|
return Ok(vec![dummy_vec; spans.len()]);
|
|
}
|
|
|
|
fn max_tokens_per_batch(&self) -> usize {
|
|
OPENAI_INPUT_LIMIT
|
|
}
|
|
|
|
fn truncate(&self, span: &str) -> (String, usize) {
|
|
let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
|
|
let token_count = tokens.len();
|
|
let output = if token_count > OPENAI_INPUT_LIMIT {
|
|
tokens.truncate(OPENAI_INPUT_LIMIT);
|
|
let new_input = OPENAI_BPE_TOKENIZER.decode(tokens.clone());
|
|
new_input.ok().unwrap_or_else(|| span.to_string())
|
|
} else {
|
|
span.to_string()
|
|
};
|
|
|
|
(output, tokens.len())
|
|
}
|
|
}
|
|
|
|
const OPENAI_INPUT_LIMIT: usize = 8190;
|
|
|
|
impl OpenAIEmbeddings {
|
|
pub fn authenticate(&mut self, cx: &mut ViewContext<Self>) {
|
|
if self.api_key.is_none() {
|
|
let api_key = if let Ok(api_key) = env::var("OPENAI_API_KEY") {
|
|
Some(api_key)
|
|
} else if let Some((_, api_key)) = cx
|
|
.platform()
|
|
.read_credentials(OPENAI_API_URL)
|
|
.log_err()
|
|
.flatten()
|
|
{
|
|
String::from_utf8(api_key).log_err()
|
|
} else {
|
|
None
|
|
};
|
|
|
|
if let Some(api_key) = api_key {
|
|
self.api_key = Some(api_key);
|
|
}
|
|
}
|
|
}
|
|
pub fn new(
|
|
api_key: Option<String>,
|
|
client: Arc<dyn HttpClient>,
|
|
executor: Arc<Background>,
|
|
) -> Self {
|
|
let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
|
|
let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
|
|
|
|
OpenAIEmbeddings {
|
|
api_key,
|
|
client,
|
|
executor,
|
|
rate_limit_count_rx,
|
|
rate_limit_count_tx,
|
|
}
|
|
}
|
|
|
|
fn resolve_rate_limit(&self) {
|
|
let reset_time = *self.rate_limit_count_tx.lock().borrow();
|
|
|
|
if let Some(reset_time) = reset_time {
|
|
if Instant::now() >= reset_time {
|
|
*self.rate_limit_count_tx.lock().borrow_mut() = None
|
|
}
|
|
}
|
|
|
|
log::trace!(
|
|
"resolving reset time: {:?}",
|
|
*self.rate_limit_count_tx.lock().borrow()
|
|
);
|
|
}
|
|
|
|
fn update_reset_time(&self, reset_time: Instant) {
|
|
let original_time = *self.rate_limit_count_tx.lock().borrow();
|
|
|
|
let updated_time = if let Some(original_time) = original_time {
|
|
if reset_time < original_time {
|
|
Some(reset_time)
|
|
} else {
|
|
Some(original_time)
|
|
}
|
|
} else {
|
|
Some(reset_time)
|
|
};
|
|
|
|
log::trace!("updating rate limit time: {:?}", updated_time);
|
|
|
|
*self.rate_limit_count_tx.lock().borrow_mut() = updated_time;
|
|
}
|
|
async fn send_request(
|
|
&self,
|
|
api_key: &str,
|
|
spans: Vec<&str>,
|
|
request_timeout: u64,
|
|
) -> Result<Response<AsyncBody>> {
|
|
let request = Request::post("https://api.openai.com/v1/embeddings")
|
|
.redirect_policy(isahc::config::RedirectPolicy::Follow)
|
|
.timeout(Duration::from_secs(request_timeout))
|
|
.header("Content-Type", "application/json")
|
|
.header("Authorization", format!("Bearer {}", api_key))
|
|
.body(
|
|
serde_json::to_string(&OpenAIEmbeddingRequest {
|
|
input: spans.clone(),
|
|
model: "text-embedding-ada-002",
|
|
})
|
|
.unwrap()
|
|
.into(),
|
|
)?;
|
|
|
|
Ok(self.client.send(request).await?)
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl EmbeddingProvider for OpenAIEmbeddings {
|
|
fn is_authenticated(&self) -> bool {
|
|
self.api_key.is_some()
|
|
}
|
|
|
|
fn max_tokens_per_batch(&self) -> usize {
|
|
50000
|
|
}
|
|
|
|
fn rate_limit_expiration(&self) -> Option<Instant> {
|
|
*self.rate_limit_count_rx.borrow()
|
|
}
|
|
fn truncate(&self, span: &str) -> (String, usize) {
|
|
let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
|
|
let output = if tokens.len() > OPENAI_INPUT_LIMIT {
|
|
tokens.truncate(OPENAI_INPUT_LIMIT);
|
|
OPENAI_BPE_TOKENIZER
|
|
.decode(tokens.clone())
|
|
.ok()
|
|
.unwrap_or_else(|| span.to_string())
|
|
} else {
|
|
span.to_string()
|
|
};
|
|
|
|
(output, tokens.len())
|
|
}
|
|
|
|
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
|
|
const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45];
|
|
const MAX_RETRIES: usize = 4;
|
|
|
|
let Some(api_key) = self.api_key.clone() else {
|
|
return Err(anyhow!("no open ai key provided"));
|
|
};
|
|
|
|
let mut request_number = 0;
|
|
let mut rate_limiting = false;
|
|
let mut request_timeout: u64 = 15;
|
|
let mut response: Response<AsyncBody>;
|
|
while request_number < MAX_RETRIES {
|
|
response = self
|
|
.send_request(
|
|
&api_key,
|
|
spans.iter().map(|x| &**x).collect(),
|
|
request_timeout,
|
|
)
|
|
.await?;
|
|
|
|
request_number += 1;
|
|
|
|
match response.status() {
|
|
StatusCode::REQUEST_TIMEOUT => {
|
|
request_timeout += 5;
|
|
}
|
|
StatusCode::OK => {
|
|
let mut body = String::new();
|
|
response.body_mut().read_to_string(&mut body).await?;
|
|
let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
|
|
|
|
log::trace!(
|
|
"openai embedding completed. tokens: {:?}",
|
|
response.usage.total_tokens
|
|
);
|
|
|
|
// If we complete a request successfully that was previously rate_limited
|
|
// resolve the rate limit
|
|
if rate_limiting {
|
|
self.resolve_rate_limit()
|
|
}
|
|
|
|
return Ok(response
|
|
.data
|
|
.into_iter()
|
|
.map(|embedding| Embedding::from(embedding.embedding))
|
|
.collect());
|
|
}
|
|
StatusCode::TOO_MANY_REQUESTS => {
|
|
rate_limiting = true;
|
|
let mut body = String::new();
|
|
response.body_mut().read_to_string(&mut body).await?;
|
|
|
|
let delay_duration = {
|
|
let delay = Duration::from_secs(BACKOFF_SECONDS[request_number - 1] as u64);
|
|
if let Some(time_to_reset) =
|
|
response.headers().get("x-ratelimit-reset-tokens")
|
|
{
|
|
if let Ok(time_str) = time_to_reset.to_str() {
|
|
parse(time_str).unwrap_or(delay)
|
|
} else {
|
|
delay
|
|
}
|
|
} else {
|
|
delay
|
|
}
|
|
};
|
|
|
|
// If we've previously rate limited, increment the duration but not the count
|
|
let reset_time = Instant::now().add(delay_duration);
|
|
self.update_reset_time(reset_time);
|
|
|
|
log::trace!(
|
|
"openai rate limiting: waiting {:?} until lifted",
|
|
&delay_duration
|
|
);
|
|
|
|
self.executor.timer(delay_duration).await;
|
|
}
|
|
_ => {
|
|
let mut body = String::new();
|
|
response.body_mut().read_to_string(&mut body).await?;
|
|
return Err(anyhow!(
|
|
"open ai bad request: {:?} {:?}",
|
|
&response.status(),
|
|
body
|
|
));
|
|
}
|
|
}
|
|
}
|
|
Err(anyhow!("openai max retries"))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rand::prelude::*;
|
|
|
|
#[gpui::test]
|
|
fn test_similarity(mut rng: StdRng) {
|
|
assert_eq!(
|
|
Embedding::from(vec![1., 0., 0., 0., 0.])
|
|
.similarity(&Embedding::from(vec![0., 1., 0., 0., 0.])),
|
|
0.
|
|
);
|
|
assert_eq!(
|
|
Embedding::from(vec![2., 0., 0., 0., 0.])
|
|
.similarity(&Embedding::from(vec![3., 1., 0., 0., 0.])),
|
|
6.
|
|
);
|
|
|
|
for _ in 0..100 {
|
|
let size = 1536;
|
|
let mut a = vec![0.; size];
|
|
let mut b = vec![0.; size];
|
|
for (a, b) in a.iter_mut().zip(b.iter_mut()) {
|
|
*a = rng.gen();
|
|
*b = rng.gen();
|
|
}
|
|
let a = Embedding::from(a);
|
|
let b = Embedding::from(b);
|
|
|
|
assert_eq!(
|
|
round_to_decimals(a.similarity(&b), 1),
|
|
round_to_decimals(reference_dot(&a.0, &b.0), 1)
|
|
);
|
|
}
|
|
|
|
fn round_to_decimals(n: OrderedFloat<f32>, decimal_places: i32) -> f32 {
|
|
let factor = (10.0 as f32).powi(decimal_places);
|
|
(n * factor).round() / factor
|
|
}
|
|
|
|
fn reference_dot(a: &[f32], b: &[f32]) -> OrderedFloat<f32> {
|
|
OrderedFloat(a.iter().zip(b.iter()).map(|(a, b)| a * b).sum())
|
|
}
|
|
}
|
|
}
|