SystemTime for Instant throughout rate_limit_expiry tracking
This commit is contained in:
KCaverly 2023-09-11 10:26:14 -04:00
parent 7df21f86dd
commit e678c7d9ee
4 changed files with 16 additions and 16 deletions

View file

@ -34,7 +34,7 @@ use std::{
ops::{Not, Range}, ops::{Not, Range},
path::PathBuf, path::PathBuf,
sync::Arc, sync::Arc,
time::{Duration, SystemTime}, time::{Duration, Instant},
}; };
use util::ResultExt as _; use util::ResultExt as _;
use workspace::{ use workspace::{
@ -329,9 +329,9 @@ impl View for ProjectSearchView {
Some(format!("Indexing...")) Some(format!("Indexing..."))
} else { } else {
if let Some(rate_limit_expiry) = rate_limit_expiry { if let Some(rate_limit_expiry) = rate_limit_expiry {
if let Ok(remaining_seconds) = let remaining_seconds =
rate_limit_expiry.duration_since(SystemTime::now()) rate_limit_expiry.duration_since(Instant::now());
{ if remaining_seconds > Duration::from_secs(0) {
Some(format!( Some(format!(
"Remaining files to index (rate limit resets in {}s): {}", "Remaining files to index (rate limit resets in {}s): {}",
remaining_seconds.as_secs(), remaining_seconds.as_secs(),

View file

@ -16,7 +16,7 @@ use serde::{Deserialize, Serialize};
use std::env; use std::env;
use std::ops::Add; use std::ops::Add;
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, SystemTime}; use std::time::{Duration, Instant};
use tiktoken_rs::{cl100k_base, CoreBPE}; use tiktoken_rs::{cl100k_base, CoreBPE};
use util::http::{HttpClient, Request}; use util::http::{HttpClient, Request};
@ -85,8 +85,8 @@ impl ToSql for Embedding {
pub struct OpenAIEmbeddings { pub struct OpenAIEmbeddings {
pub client: Arc<dyn HttpClient>, pub client: Arc<dyn HttpClient>,
pub executor: Arc<Background>, pub executor: Arc<Background>,
rate_limit_count_rx: watch::Receiver<Option<SystemTime>>, rate_limit_count_rx: watch::Receiver<Option<Instant>>,
rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<SystemTime>>>>, rate_limit_count_tx: Arc<Mutex<watch::Sender<Option<Instant>>>>,
} }
#[derive(Serialize)] #[derive(Serialize)]
@ -119,14 +119,14 @@ pub trait EmbeddingProvider: Sync + Send {
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>; async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
fn max_tokens_per_batch(&self) -> usize; fn max_tokens_per_batch(&self) -> usize;
fn truncate(&self, span: &str) -> (String, usize); fn truncate(&self, span: &str) -> (String, usize);
fn rate_limit_expiration(&self) -> Option<SystemTime>; fn rate_limit_expiration(&self) -> Option<Instant>;
} }
pub struct DummyEmbeddings {} pub struct DummyEmbeddings {}
#[async_trait] #[async_trait]
impl EmbeddingProvider for DummyEmbeddings { impl EmbeddingProvider for DummyEmbeddings {
fn rate_limit_expiration(&self) -> Option<SystemTime> { fn rate_limit_expiration(&self) -> Option<Instant> {
None None
} }
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> { async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
@ -174,7 +174,7 @@ impl OpenAIEmbeddings {
let reset_time = *self.rate_limit_count_tx.lock().borrow(); let reset_time = *self.rate_limit_count_tx.lock().borrow();
if let Some(reset_time) = reset_time { if let Some(reset_time) = reset_time {
if SystemTime::now() >= reset_time { if Instant::now() >= reset_time {
*self.rate_limit_count_tx.lock().borrow_mut() = None *self.rate_limit_count_tx.lock().borrow_mut() = None
} }
} }
@ -185,7 +185,7 @@ impl OpenAIEmbeddings {
); );
} }
fn update_reset_time(&self, reset_time: SystemTime) { fn update_reset_time(&self, reset_time: Instant) {
let original_time = *self.rate_limit_count_tx.lock().borrow(); let original_time = *self.rate_limit_count_tx.lock().borrow();
let updated_time = if let Some(original_time) = original_time { let updated_time = if let Some(original_time) = original_time {
@ -232,7 +232,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
50000 50000
} }
fn rate_limit_expiration(&self) -> Option<SystemTime> { fn rate_limit_expiration(&self) -> Option<Instant> {
*self.rate_limit_count_rx.borrow() *self.rate_limit_count_rx.borrow()
} }
fn truncate(&self, span: &str) -> (String, usize) { fn truncate(&self, span: &str) -> (String, usize) {
@ -319,7 +319,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
}; };
// If we've previously rate limited, increment the duration but not the count // If we've previously rate limited, increment the duration but not the count
let reset_time = SystemTime::now().add(delay_duration); let reset_time = Instant::now().add(delay_duration);
self.update_reset_time(reset_time); self.update_reset_time(reset_time);
log::trace!( log::trace!(

View file

@ -112,7 +112,7 @@ pub enum SemanticIndexStatus {
Indexed, Indexed,
Indexing { Indexing {
remaining_files: usize, remaining_files: usize,
rate_limit_expiry: Option<SystemTime>, rate_limit_expiry: Option<Instant>,
}, },
} }

View file

@ -21,7 +21,7 @@ use std::{
atomic::{self, AtomicUsize}, atomic::{self, AtomicUsize},
Arc, Arc,
}, },
time::SystemTime, time::{Instant, SystemTime},
}; };
use unindent::Unindent; use unindent::Unindent;
use util::RandomCharIter; use util::RandomCharIter;
@ -1275,7 +1275,7 @@ impl EmbeddingProvider for FakeEmbeddingProvider {
200 200
} }
fn rate_limit_expiration(&self) -> Option<SystemTime> { fn rate_limit_expiration(&self) -> Option<Instant> {
None None
} }