catchup with main

crates/semantic_index/src/semantic_index.rs

@@ -124,6 +124,8 @@ pub struct SemanticIndex {
     _embedding_task: Task<()>,
     _parsing_files_tasks: Vec<Task<()>>,
     projects: HashMap<WeakModelHandle<Project>, ProjectState>,
+    api_key: Option<String>,
+    embedding_queue: Arc<Mutex<EmbeddingQueue>>,
 }
 
 struct ProjectState {
@@ -269,7 +271,7 @@ pub struct SearchResult {
 }
 
 impl SemanticIndex {
-    pub fn global(cx: &AppContext) -> Option<ModelHandle<SemanticIndex>> {
+    pub fn global(cx: &mut AppContext) -> Option<ModelHandle<SemanticIndex>> {
         if cx.has_global::<ModelHandle<Self>>() {
             Some(cx.global::<ModelHandle<SemanticIndex>>().clone())
         } else {
@@ -277,12 +279,26 @@ impl SemanticIndex {
         }
     }
 
+    pub fn authenticate(&mut self, cx: &AppContext) {
+        if self.api_key.is_none() {
+            self.api_key = self.embedding_provider.retrieve_credentials(cx);
+
+            self.embedding_queue
+                .lock()
+                .set_api_key(self.api_key.clone());
+        }
+    }
+
+    pub fn is_authenticated(&self) -> bool {
+        self.api_key.is_some()
+    }
+
     pub fn enabled(cx: &AppContext) -> bool {
         settings::get::<SemanticIndexSettings>(cx).enabled
     }
 
     pub fn status(&self, project: &ModelHandle<Project>) -> SemanticIndexStatus {
-        if !self.embedding_provider.is_authenticated() {
+        if !self.is_authenticated() {
             return SemanticIndexStatus::NotAuthenticated;
         }
 
@@ -324,7 +340,7 @@ impl SemanticIndex {
         Ok(cx.add_model(|cx| {
             let t0 = Instant::now();
             let embedding_queue =
-                EmbeddingQueue::new(embedding_provider.clone(), cx.background().clone());
+                EmbeddingQueue::new(embedding_provider.clone(), cx.background().clone(), None);
             let _embedding_task = cx.background().spawn({
                 let embedded_files = embedding_queue.finished_files();
                 let db = db.clone();
@@ -389,6 +405,8 @@ impl SemanticIndex {
                 _embedding_task,
                 _parsing_files_tasks,
                 projects: Default::default(),
+                api_key: None,
+                embedding_queue
             }
         }))
     }
@@ -703,12 +721,13 @@ impl SemanticIndex {
 
         let index = self.index_project(project.clone(), cx);
         let embedding_provider = self.embedding_provider.clone();
+        let api_key = self.api_key.clone();
 
         cx.spawn(|this, mut cx| async move {
             index.await?;
             let t0 = Instant::now();
             let query = embedding_provider
-                .embed_batch(vec![query])
+                .embed_batch(vec![query], api_key)
                 .await?
                 .pop()
                 .ok_or_else(|| anyhow!("could not embed query"))?;
@@ -926,6 +945,7 @@ impl SemanticIndex {
         let fs = self.fs.clone();
         let db_path = self.db.path().clone();
         let background = cx.background().clone();
+        let api_key = self.api_key.clone();
         cx.background().spawn(async move {
             let db = VectorDatabase::new(fs, db_path.clone(), background).await?;
             let mut results = Vec::<SearchResult>::new();
@@ -940,10 +960,15 @@ impl SemanticIndex {
                     .parse_file_with_template(None, &snapshot.text(), language)
                     .log_err()
                     .unwrap_or_default();
-                if Self::embed_spans(&mut spans, embedding_provider.as_ref(), &db)
-                    .await
-                    .log_err()
-                    .is_some()
+                if Self::embed_spans(
+                    &mut spans,
+                    embedding_provider.as_ref(),
+                    &db,
+                    api_key.clone(),
+                )
+                .await
+                .log_err()
+                .is_some()
                 {
                     for span in spans {
                         let similarity = span.embedding.unwrap().similarity(&query);
@@ -983,8 +1008,11 @@ impl SemanticIndex {
         project: ModelHandle<Project>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<()>> {
-        if !self.embedding_provider.is_authenticated() {
-            return Task::ready(Err(anyhow!("user is not authenticated")));
+        if self.api_key.is_none() {
+            self.authenticate(cx);
+            if self.api_key.is_none() {
+                return Task::ready(Err(anyhow!("user is not authenticated")));
+            }
         }
 
         if !self.projects.contains_key(&project.downgrade()) {
@@ -1165,6 +1193,7 @@ impl SemanticIndex {
         spans: &mut [Span],
         embedding_provider: &dyn EmbeddingProvider,
         db: &VectorDatabase,
+        api_key: Option<String>,
     ) -> Result<()> {
         let mut batch = Vec::new();
         let mut batch_tokens = 0;
@@ -1187,7 +1216,7 @@ impl SemanticIndex {
 
             if batch_tokens + span.token_count > embedding_provider.max_tokens_per_batch() {
                 let batch_embeddings = embedding_provider
-                    .embed_batch(mem::take(&mut batch))
+                    .embed_batch(mem::take(&mut batch), api_key.clone())
                     .await?;
                 embeddings.extend(batch_embeddings);
                 batch_tokens = 0;
@@ -1199,7 +1228,7 @@ impl SemanticIndex {
 
         if !batch.is_empty() {
             let batch_embeddings = embedding_provider
-                .embed_batch(mem::take(&mut batch))
+                .embed_batch(mem::take(&mut batch), api_key)
                 .await?;
 
             embeddings.extend(batch_embeddings);