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

moved authentication for the semantic index into the EmbeddingProvider

Browse source
This commit is contained in:
KCaverly 2023-10-30 10:02:27 -04:00
parent 1e8b23d8fb
commit a2c3971ad6
14 changed files with 200 additions and 206 deletions
Download patch file Download diff file Expand all files Collapse all files

17
crates/semantic_index/src/embedding_queue.rs
View file

@ -1,5 +1,5 @@
use crate::{parsing::Span, JobHandle};
use ai::{auth::ProviderCredential, embedding::EmbeddingProvider};
use ai::embedding::EmbeddingProvider;
use gpui::executor::Background;
use parking_lot::Mutex;
use smol::channel;
@ -41,7 +41,6 @@ pub struct EmbeddingQueue {
pending_batch_token_count: usize,
finished_files_tx: channel::Sender<FileToEmbed>,
finished_files_rx: channel::Receiver<FileToEmbed>,
pub provider_credential: ProviderCredential,
}
#[derive(Clone)]
@ -51,11 +50,7 @@ pub struct FileFragmentToEmbed {
}
impl EmbeddingQueue {
pub fn new(
embedding_provider: Arc<dyn EmbeddingProvider>,
executor: Arc<Background>,
provider_credential: ProviderCredential,
) -> Self {
pub fn new(embedding_provider: Arc<dyn EmbeddingProvider>, executor: Arc<Background>) -> Self {
let (finished_files_tx, finished_files_rx) = channel::unbounded();
Self {
embedding_provider,
@ -64,14 +59,9 @@ impl EmbeddingQueue {
pending_batch_token_count: 0,
finished_files_tx,
finished_files_rx,
provider_credential,
}
}
pub fn set_credential(&mut self, credential: ProviderCredential) {
self.provider_credential = credential;
}
pub fn push(&mut self, file: FileToEmbed) {
if file.spans.is_empty() {
self.finished_files_tx.try_send(file).unwrap();
@ -118,7 +108,6 @@ impl EmbeddingQueue {
let finished_files_tx = self.finished_files_tx.clone();
let embedding_provider = self.embedding_provider.clone();
let credential = self.provider_credential.clone();
self.executor
.spawn(async move {
@ -143,7 +132,7 @@ impl EmbeddingQueue {
return;
};
match embedding_provider.embed_batch(spans, credential).await {
match embedding_provider.embed_batch(spans).await {
Ok(embeddings) => {
let mut embeddings = embeddings.into_iter();
for fragment in batch {