Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

updated semantic_index reset status to leverage target reset system time as opposed to duration

Browse source
This commit is contained in:
KCaverly 2023-09-08 15:04:50 -04:00
parent a5ee8fc805
commit bf43f93197
4 changed files with 50 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

67
crates/semantic_index/src/embedding.rs
View file

@ -14,8 +14,9 @@ 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;
use std::time::{Duration, SystemTime};
use tiktoken_rs::{cl100k_base, CoreBPE};
use util::http::{HttpClient, Request};
@ -84,8 +85,8 @@ impl ToSql for Embedding {
pub struct OpenAIEmbeddings {
pub client: Arc<dyn HttpClient>,
pub executor: Arc<Background>,
rate_limit_count_rx: watch::Receiver<(Duration, usize)>,
rate_limit_count_tx: Arc<Mutex<watch::Sender<(Duration, usize)>>>,
rate_limit_count_rx: watch::Receiver<(Option<SystemTime>, usize)>,
rate_limit_count_tx: Arc<Mutex<watch::Sender<(Option<SystemTime>, usize)>>>,
}
#[derive(Serialize)]
@ -118,15 +119,15 @@ pub trait EmbeddingProvider: Sync + Send {
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) -> Duration;
fn rate_limit_expiration(&self) -> Option<SystemTime>;
}
pub struct DummyEmbeddings {}
#[async_trait]
impl EmbeddingProvider for DummyEmbeddings {
fn rate_limit_expiration(&self) -> Duration {
Duration::ZERO
fn rate_limit_expiration(&self) -> Option<SystemTime> {
None
}
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
// 1024 is the OpenAI Embeddings size for ada models.
@ -158,7 +159,7 @@ const OPENAI_INPUT_LIMIT: usize = 8190;
impl OpenAIEmbeddings {
pub fn new(client: Arc<dyn HttpClient>, executor: Arc<Background>) -> Self {
let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with((Duration::ZERO, 0));
let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with((None, 0));
let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
OpenAIEmbeddings {
@ -170,39 +171,32 @@ impl OpenAIEmbeddings {
}
fn resolve_rate_limit(&self) {
let (current_delay, delay_count) = *self.rate_limit_count_tx.lock().borrow();
let (reset_time, delay_count) = *self.rate_limit_count_tx.lock().borrow();
let updated_count = delay_count - 1;
let updated_duration = if updated_count == 0 {
Duration::ZERO
} else {
current_delay
};
let updated_time = if updated_count == 0 { None } else { reset_time };
log::trace!(
"resolving rate limit: Count: {:?} Duration: {:?}",
updated_count,
updated_duration
);
log::trace!("resolving rate limit: Count: {:?}", updated_count);
*self.rate_limit_count_tx.lock().borrow_mut() = (updated_duration, updated_count);
*self.rate_limit_count_tx.lock().borrow_mut() = (updated_time, updated_count);
}
fn update_rate_limit(&self, delay_duration: Duration, count_increase: usize) {
let (current_delay, delay_count) = *self.rate_limit_count_tx.lock().borrow();
let updated_count = delay_count + count_increase;
let updated_duration = if current_delay < delay_duration {
delay_duration
fn update_rate_limit(&self, reset_time: SystemTime, count_increase: usize) {
let (original_time, original_count) = *self.rate_limit_count_tx.lock().borrow();
let updated_count = original_count + count_increase;
let updated_time = if let Some(original_time) = original_time {
if reset_time < original_time {
Some(reset_time)
} else {
Some(original_time)
}
} else {
current_delay
Some(reset_time)
};
log::trace!(
"updating rate limit: Count: {:?} Duration: {:?}",
updated_count,
updated_duration
);
log::trace!("updating rate limit: Count: {:?}", updated_count);
*self.rate_limit_count_tx.lock().borrow_mut() = (updated_duration, updated_count);
*self.rate_limit_count_tx.lock().borrow_mut() = (updated_time, updated_count);
}
async fn send_request(
&self,
@ -234,9 +228,9 @@ impl EmbeddingProvider for OpenAIEmbeddings {
50000
}
fn rate_limit_expiration(&self) -> Duration {
let (duration, _) = *self.rate_limit_count_rx.borrow();
duration
fn rate_limit_expiration(&self) -> Option<SystemTime> {
let (expiration_time, _) = *self.rate_limit_count_rx.borrow();
expiration_time
}
fn truncate(&self, span: &str) -> (String, usize) {
let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
@ -321,10 +315,11 @@ impl EmbeddingProvider for OpenAIEmbeddings {
};
// If we've previously rate limited, increment the duration but not the count
let reset_time = SystemTime::now().add(delay_duration);
if rate_limiting {
self.update_rate_limit(delay_duration, 0);
self.update_rate_limit(reset_time, 0);
} else {
self.update_rate_limit(delay_duration, 1);
self.update_rate_limit(reset_time, 1);
}
rate_limiting = true;

10
crates/semantic_index/src/semantic_index.rs
View file

@ -112,7 +112,7 @@ pub enum SemanticIndexStatus {
Indexed,
Indexing {
remaining_files: usize,
rate_limiting: Duration,
rate_limit_expiration_time: Option<SystemTime>,
},
}
@ -132,8 +132,6 @@ struct ProjectState {
pending_file_count_rx: watch::Receiver<usize>,
pending_file_count_tx: Arc<Mutex<watch::Sender<usize>>>,
pending_index: usize,
rate_limiting_count_rx: watch::Receiver<usize>,
rate_limiting_count_tx: Arc<Mutex<watch::Sender<usize>>>,
_subscription: gpui::Subscription,
_observe_pending_file_count: Task<()>,
}
@ -225,15 +223,11 @@ impl ProjectState {
fn new(subscription: gpui::Subscription, cx: &mut ModelContext<SemanticIndex>) -> Self {
let (pending_file_count_tx, pending_file_count_rx) = watch::channel_with(0);
let pending_file_count_tx = Arc::new(Mutex::new(pending_file_count_tx));
let (rate_limiting_count_tx, rate_limiting_count_rx) = watch::channel_with(0);
let rate_limiting_count_tx = Arc::new(Mutex::new(rate_limiting_count_tx));
Self {
worktrees: Default::default(),
pending_file_count_rx: pending_file_count_rx.clone(),
pending_file_count_tx,
pending_index: 0,
rate_limiting_count_rx: rate_limiting_count_rx.clone(),
rate_limiting_count_tx,
_subscription: subscription,
_observe_pending_file_count: cx.spawn_weak({
let mut pending_file_count_rx = pending_file_count_rx.clone();
@ -299,7 +293,7 @@ impl SemanticIndex {
} else {
SemanticIndexStatus::Indexing {
remaining_files: project_state.pending_file_count_rx.borrow().clone(),
rate_limiting: self.embedding_provider.rate_limit_expiration(),
rate_limit_expiration_time: self.embedding_provider.rate_limit_expiration(),
}
}
} else {

6
crates/semantic_index/src/semantic_index_tests.rs
View file

@ -21,7 +21,7 @@ use std::{
atomic::{self, AtomicUsize},
Arc,
},
time::{Duration, SystemTime},
time::SystemTime,
};
use unindent::Unindent;
use util::RandomCharIter;
@ -1275,8 +1275,8 @@ impl EmbeddingProvider for FakeEmbeddingProvider {
200
}
fn rate_limit_expiration(&self) -> Duration {
Duration::ZERO
fn rate_limit_expiration(&self) -> Option<SystemTime> {
None
}
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {