zed-industries/zed
1
1
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2024-11-08 07:35:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

update logging for open ai embedding and remove redundant truncation

Browse Source
This commit is contained in:
KCaverly 2023-07-18 11:00:21 -04:00
parent b9fdfd60f0
commit ed1b1a5ccd
1 changed files with 21 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
crates/semantic_index/src/embedding.rs
View File

@ -67,13 +67,17 @@ impl EmbeddingProvider for DummyEmbeddings {
}
}
const INPUT_LIMIT: usize = 8190;
const OPENAI_INPUT_LIMIT: usize = 8190;
impl OpenAIEmbeddings {
pub fn new(client: Arc<dyn HttpClient>, executor: Arc<Background>) -> Self {
Self { client, executor }
}
fn truncate(span: String) -> String {
let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span.as_ref());
if tokens.len() > INPUT_LIMIT {
tokens.truncate(INPUT_LIMIT);
if tokens.len() > OPENAI_INPUT_LIMIT {
tokens.truncate(OPENAI_INPUT_LIMIT);
let result = OPENAI_BPE_TOKENIZER.decode(tokens.clone());
if result.is_ok() {
let transformed = result.unwrap();
@ -115,6 +119,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
.ok_or_else(|| anyhow!("no api key"))?;
let mut request_number = 0;
let mut truncated = false;
let mut response: Response<AsyncBody>;
let mut spans: Vec<String> = spans.iter().map(|x| x.to_string()).collect();
while request_number < MAX_RETRIES {
@ -136,15 +141,18 @@ impl EmbeddingProvider for OpenAIEmbeddings {
self.executor.timer(delay).await;
}
StatusCode::BAD_REQUEST => {
log::info!(
"BAD REQUEST: {:?} {:?}",
&response.status(),
response.body()
);
// Don't worry about delaying bad request, as we can assume
// we haven't been rate limited yet.
// Only truncate if it hasnt been truncated before
if !truncated {
for span in spans.iter_mut() {
*span = Self::truncate(span.to_string());
*span = Self::truncate(span.clone());
}
truncated = true;
} else {
// If failing once already truncated, log the error and break the loop
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
log::trace!("open ai bad request: {:?} {:?}", &response.status(), body);
break;
}
}
StatusCode::OK => {
@ -152,7 +160,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
response.body_mut().read_to_string(&mut body).await?;
let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
log::info!(
log::trace!(
"openai embedding completed. tokens: {:?}",
response.usage.total_tokens
);