1use anyhow::{anyhow, Result};
2use async_trait::async_trait;
3use futures::AsyncReadExt;
4use gpui::executor::Background;
5use gpui::serde_json;
6use isahc::http::StatusCode;
7use isahc::prelude::Configurable;
8use isahc::{AsyncBody, Response};
9use lazy_static::lazy_static;
10use ordered_float::OrderedFloat;
11use parking_lot::Mutex;
12use parse_duration::parse;
13use postage::watch;
14use rusqlite::types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef};
15use rusqlite::ToSql;
16use serde::{Deserialize, Serialize};
17use std::env;
18use std::ops::Add;
19use std::sync::Arc;
20use std::time::{Duration, Instant};
21use tiktoken_rs::{cl100k_base, CoreBPE};
22use util::http::{HttpClient, Request};
23
24lazy_static! {
25 static ref OPENAI_API_KEY: Option<String> = env::var("OPENAI_API_KEY").ok();
26 static ref OPENAI_BPE_TOKENIZER: CoreBPE = cl100k_base().unwrap();
27}
28
29#[derive(Debug, PartialEq, Clone)]
30pub struct Embedding(Vec<f32>);
31
32impl From<Vec<f32>> for Embedding {
33 fn from(value: Vec<f32>) -> Self {
34 Embedding(value)
35 }
36}
37
38impl Embedding {
39 pub fn similarity(&self, other: &Self) -> OrderedFloat<f32> {
40 let len = self.0.len();
41 assert_eq!(len, other.0.len());
42
43 let mut result = 0.0;
44 unsafe {
45 matrixmultiply::sgemm(
46 1,
47 len,
48 1,
49 1.0,
50 self.0.as_ptr(),
51 len as isize,
52 1,
53 other.0.as_ptr(),
54 1,
55 len as isize,
56 0.0,
57 &mut result as *mut f32,
58 1,
59 1,
60 );
61 }
62 OrderedFloat(result)
63 }
64}
65
66impl FromSql for Embedding {
67 fn column_result(value: ValueRef) -> FromSqlResult<Self> {
68 let bytes = value.as_blob()?;
69 let embedding: Result<Vec<f32>, Box<bincode::ErrorKind>> = bincode::deserialize(bytes);
70 if embedding.is_err() {
71 return Err(rusqlite::types::FromSqlError::Other(embedding.unwrap_err()));
72 }
73 Ok(Embedding(embedding.unwrap()))
74 }
75}
76
77impl ToSql for Embedding {
78 fn to_sql(&self) -> rusqlite::Result<ToSqlOutput> {
79 let bytes = bincode::serialize(&self.0)
80 .map_err(|err| rusqlite::Error::ToSqlConversionFailure(Box::new(err)))?;
81 Ok(ToSqlOutput::Owned(rusqlite::types::Value::Blob(bytes)))
82 }
83}
84
85#[derive(Clone)]
86pub struct OpenAIEmbeddings {
87 pub client: Arc<dyn HttpClient>,
88 pub executor: Arc<Background>,
89 rate_limit_count_rx: watch::Receiver<Option<Instant>>,
90 rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
91}
92
93#[derive(Serialize)]
94struct OpenAIEmbeddingRequest<'a> {
95 model: &'static str,
96 input: Vec<&'a str>,
97}
98
99#[derive(Deserialize)]
100struct OpenAIEmbeddingResponse {
101 data: Vec<OpenAIEmbedding>,
102 usage: OpenAIEmbeddingUsage,
103}
104
105#[derive(Debug, Deserialize)]
106struct OpenAIEmbedding {
107 embedding: Vec<f32>,
108 index: usize,
109 object: String,
110}
111
112#[derive(Deserialize)]
113struct OpenAIEmbeddingUsage {
114 prompt_tokens: usize,
115 total_tokens: usize,
116}
117
118#[async_trait]
119pub trait EmbeddingProvider: Sync + Send {
120 async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
121 fn max_tokens_per_batch(&self) -> usize;
122 fn truncate(&self, span: &str) -> (String, usize);
123 fn rate_limit_expiration(&self) -> Option<Instant>;
124}
125
126pub struct DummyEmbeddings {}
127
128#[async_trait]
129impl EmbeddingProvider for DummyEmbeddings {
130 fn rate_limit_expiration(&self) -> Option<Instant> {
131 None
132 }
133 async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
134 // 1024 is the OpenAI Embeddings size for ada models.
135 // the model we will likely be starting with.
136 let dummy_vec = Embedding::from(vec![0.32 as f32; 1536]);
137 return Ok(vec![dummy_vec; spans.len()]);
138 }
139
140 fn max_tokens_per_batch(&self) -> usize {
141 OPENAI_INPUT_LIMIT
142 }
143
144 fn truncate(&self, span: &str) -> (String, usize) {
145 let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
146 let token_count = tokens.len();
147 let output = if token_count > OPENAI_INPUT_LIMIT {
148 tokens.truncate(OPENAI_INPUT_LIMIT);
149 let new_input = OPENAI_BPE_TOKENIZER.decode(tokens.clone());
150 new_input.ok().unwrap_or_else(|| span.to_string())
151 } else {
152 span.to_string()
153 };
154
155 (output, tokens.len())
156 }
157}
158
159const OPENAI_INPUT_LIMIT: usize = 8190;
160
161impl OpenAIEmbeddings {
162 pub fn new(client: Arc<dyn HttpClient>, executor: Arc<Background>) -> Self {
163 let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
164 let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
165
166 OpenAIEmbeddings {
167 client,
168 executor,
169 rate_limit_count_rx,
170 rate_limit_count_tx,
171 }
172 }
173
174 fn resolve_rate_limit(&self) {
175 let reset_time = *self.rate_limit_count_tx.lock().borrow();
176
177 if let Some(reset_time) = reset_time {
178 if Instant::now() >= reset_time {
179 *self.rate_limit_count_tx.lock().borrow_mut() = None
180 }
181 }
182
183 log::trace!(
184 "resolving reset time: {:?}",
185 *self.rate_limit_count_tx.lock().borrow()
186 );
187 }
188
189 fn update_reset_time(&self, reset_time: Instant) {
190 let original_time = *self.rate_limit_count_tx.lock().borrow();
191
192 let updated_time = if let Some(original_time) = original_time {
193 if reset_time < original_time {
194 Some(reset_time)
195 } else {
196 Some(original_time)
197 }
198 } else {
199 Some(reset_time)
200 };
201
202 log::trace!("updating rate limit time: {:?}", updated_time);
203
204 *self.rate_limit_count_tx.lock().borrow_mut() = updated_time;
205 }
206 async fn send_request(
207 &self,
208 api_key: &str,
209 spans: Vec<&str>,
210 request_timeout: u64,
211 ) -> Result<Response<AsyncBody>> {
212 let request = Request::post("https://api.openai.com/v1/embeddings")
213 .redirect_policy(isahc::config::RedirectPolicy::Follow)
214 .timeout(Duration::from_secs(request_timeout))
215 .header("Content-Type", "application/json")
216 .header("Authorization", format!("Bearer {}", api_key))
217 .body(
218 serde_json::to_string(&OpenAIEmbeddingRequest {
219 input: spans.clone(),
220 model: "text-embedding-ada-002",
221 })
222 .unwrap()
223 .into(),
224 )?;
225
226 Ok(self.client.send(request).await?)
227 }
228}
229
230#[async_trait]
231impl EmbeddingProvider for OpenAIEmbeddings {
232 fn max_tokens_per_batch(&self) -> usize {
233 50000
234 }
235
236 fn rate_limit_expiration(&self) -> Option<Instant> {
237 *self.rate_limit_count_rx.borrow()
238 }
239 fn truncate(&self, span: &str) -> (String, usize) {
240 let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
241 let output = if tokens.len() > OPENAI_INPUT_LIMIT {
242 tokens.truncate(OPENAI_INPUT_LIMIT);
243 OPENAI_BPE_TOKENIZER
244 .decode(tokens.clone())
245 .ok()
246 .unwrap_or_else(|| span.to_string())
247 } else {
248 span.to_string()
249 };
250
251 (output, tokens.len())
252 }
253
254 async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
255 const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45];
256 const MAX_RETRIES: usize = 4;
257
258 let api_key = OPENAI_API_KEY
259 .as_ref()
260 .ok_or_else(|| anyhow!("no api key"))?;
261
262 let mut request_number = 0;
263 let mut rate_limiting = false;
264 let mut request_timeout: u64 = 15;
265 let mut response: Response<AsyncBody>;
266 while request_number < MAX_RETRIES {
267 response = self
268 .send_request(
269 api_key,
270 spans.iter().map(|x| &**x).collect(),
271 request_timeout,
272 )
273 .await?;
274 request_number += 1;
275
276 match response.status() {
277 StatusCode::REQUEST_TIMEOUT => {
278 request_timeout += 5;
279 }
280 StatusCode::OK => {
281 let mut body = String::new();
282 response.body_mut().read_to_string(&mut body).await?;
283 let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
284
285 log::trace!(
286 "openai embedding completed. tokens: {:?}",
287 response.usage.total_tokens
288 );
289
290 // If we complete a request successfully that was previously rate_limited
291 // resolve the rate limit
292 if rate_limiting {
293 self.resolve_rate_limit()
294 }
295
296 return Ok(response
297 .data
298 .into_iter()
299 .map(|embedding| Embedding::from(embedding.embedding))
300 .collect());
301 }
302 StatusCode::TOO_MANY_REQUESTS => {
303 rate_limiting = true;
304 let mut body = String::new();
305 response.body_mut().read_to_string(&mut body).await?;
306
307 let delay_duration = {
308 let delay = Duration::from_secs(BACKOFF_SECONDS[request_number - 1] as u64);
309 if let Some(time_to_reset) =
310 response.headers().get("x-ratelimit-reset-tokens")
311 {
312 if let Ok(time_str) = time_to_reset.to_str() {
313 parse(time_str).unwrap_or(delay)
314 } else {
315 delay
316 }
317 } else {
318 delay
319 }
320 };
321
322 // If we've previously rate limited, increment the duration but not the count
323 let reset_time = Instant::now().add(delay_duration);
324 self.update_reset_time(reset_time);
325
326 log::trace!(
327 "openai rate limiting: waiting {:?} until lifted",
328 &delay_duration
329 );
330
331 self.executor.timer(delay_duration).await;
332 }
333 _ => {
334 let mut body = String::new();
335 response.body_mut().read_to_string(&mut body).await?;
336 return Err(anyhow!(
337 "open ai bad request: {:?} {:?}",
338 &response.status(),
339 body
340 ));
341 }
342 }
343 }
344 Err(anyhow!("openai max retries"))
345 }
346}
347
348#[cfg(test)]
349mod tests {
350 use super::*;
351 use rand::prelude::*;
352
353 #[gpui::test]
354 fn test_similarity(mut rng: StdRng) {
355 assert_eq!(
356 Embedding::from(vec![1., 0., 0., 0., 0.])
357 .similarity(&Embedding::from(vec![0., 1., 0., 0., 0.])),
358 0.
359 );
360 assert_eq!(
361 Embedding::from(vec![2., 0., 0., 0., 0.])
362 .similarity(&Embedding::from(vec![3., 1., 0., 0., 0.])),
363 6.
364 );
365
366 for _ in 0..100 {
367 let size = 1536;
368 let mut a = vec![0.; size];
369 let mut b = vec![0.; size];
370 for (a, b) in a.iter_mut().zip(b.iter_mut()) {
371 *a = rng.gen();
372 *b = rng.gen();
373 }
374 let a = Embedding::from(a);
375 let b = Embedding::from(b);
376
377 assert_eq!(
378 round_to_decimals(a.similarity(&b), 1),
379 round_to_decimals(reference_dot(&a.0, &b.0), 1)
380 );
381 }
382
383 fn round_to_decimals(n: OrderedFloat<f32>, decimal_places: i32) -> f32 {
384 let factor = (10.0 as f32).powi(decimal_places);
385 (n * factor).round() / factor
386 }
387
388 fn reference_dot(a: &[f32], b: &[f32]) -> OrderedFloat<f32> {
389 OrderedFloat(a.iter().zip(b.iter()).map(|(a, b)| a * b).sum())
390 }
391 }
392}