@@ -34,7 +34,7 @@ use std::{
ops::{Not, Range},
path::PathBuf,
sync::Arc,
- time::Duration,
+ time::SystemTime,
};
use util::ResultExt as _;
use workspace::{
@@ -322,17 +322,23 @@ impl View for ProjectSearchView {
SemanticIndexStatus::Indexed => Some("Indexing complete".to_string()),
SemanticIndexStatus::Indexing {
remaining_files,
- rate_limiting,
+ rate_limit_expiration_time,
} => {
if remaining_files == 0 {
Some(format!("Indexing..."))
} else {
- if rate_limiting > Duration::ZERO {
- Some(format!(
- "Remaining files to index (rate limit resets in {}s): {}",
- rate_limiting.as_secs(),
- remaining_files
- ))
+ if let Some(rate_limit_expiration_time) = rate_limit_expiration_time {
+ if let Ok(remaining_seconds) =
+ rate_limit_expiration_time.duration_since(SystemTime::now())
+ {
+ Some(format!(
+ "Remaining files to index(rate limit resets in {}s): {}",
+ remaining_seconds.as_secs(),
+ remaining_files
+ ))
+ } else {
+ Some(format!("Remaining files to index: {}", remaining_files))
+ }
} else {
Some(format!("Remaining files to index: {}", remaining_files))
}
@@ -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;
@@ -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 {
@@ -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>> {