improve perf

Author: thesandybridge

Cargo.lock

@@ -27,3 +27,4 @@ bat = { version = "0.25.0", features = ["build-assets"] }
 ignore = "0.4.23"
 sha2 = "0.10.9"
 hex = "0.4.3"
+futures = "0.3.31"

kb_core/Cargo.toml

@@ -3,13 +3,18 @@ use crate::config;
 use crate::embedding;
 use crate::llm;
 use crate::utils;
-use std::time::UNIX_EPOCH;
+use crate::state::{IndexState, IndexedChunk};
+use futures::stream::{FuturesUnordered, StreamExt};
+use std::time::{Duration, UNIX_EPOCH};
+use tokio::time::sleep;
 use indicatif::{ProgressBar, ProgressStyle};
 use reqwest::Client;
 use std::fs;
 use std::path::Path;
 use uuid::Uuid;
 
+const BATCH_SIZE: usize = 8;
+
 pub fn handle_config(set_api_key: Option<String>, show: bool) -> anyhow::Result<()> {
     let config_path = dirs::config_dir()
         .ok_or_else(|| anyhow::anyhow!("Unable to determine config directory"))?
@@ -57,7 +62,7 @@ pub fn handle_config(set_api_key: Option<String>, show: bool) -> anyhow::Result<
 }
 
 pub async fn handle_index(client: &Client, path: &Path) -> anyhow::Result<()> {
-    let paths = crate::utils::collect_files(path)?;
+    let paths = utils::collect_files(path)?;
     let total_files = paths.len() as u64;
     let pb = ProgressBar::new(total_files);
     pb.set_style(
@@ -66,16 +71,16 @@ pub async fn handle_index(client: &Client, path: &Path) -> anyhow::Result<()> {
             .tick_chars("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏ "),
     );
 
-    let config_dir = crate::config::get_config_dir()?;
-    let mut state = crate::state::IndexState::load(&config_dir)?;
+    let config_dir = config::get_config_dir()?;
+    let mut state = IndexState::load(&config_dir)?;
 
     for path in paths {
         pb.set_message(format!("Indexing {}", path.display()));
         let metadata = fs::metadata(&path)?;
         let modified = metadata.modified()?.duration_since(UNIX_EPOCH)?.as_secs();
         let file_str = path.to_string_lossy().to_string();
 
-        // Skip if the file hasn't changed
+        // Skip if file unchanged
         if let Some(prev) = state.get_last_modified(&file_str) {
             if prev == modified {
                 pb.inc(1);
@@ -84,30 +89,50 @@ pub async fn handle_index(client: &Client, path: &Path) -> anyhow::Result<()> {
         }
 
         let content = fs::read_to_string(&path)?;
-        let chunks = crate::utils::chunk_text(&content);
+        let chunks = utils::chunk_text(&content);
         let prev_chunks = state.get_file_chunks(&file_str).cloned().unwrap_or_default();
-        let mut new_chunks = vec![];
+        let mut new_chunks = Vec::new();
+        let mut chunk_info = Vec::new();
 
         for chunk in &chunks {
             if chunk.trim().is_empty() || chunk.len() > 100_000 {
                 continue;
             }
 
-            let hash = crate::state::IndexState::hash_chunk(chunk);
-            if crate::state::IndexState::has_chunk(&prev_chunks, &hash) {
+            let hash = IndexState::hash_chunk(chunk);
+            if IndexState::has_chunk(&prev_chunks, &hash) {
                 continue;
             }
 
-            let id = Uuid::new_v4().to_string();
-            let embedding = crate::embedding::get_embedding(client, chunk).await?;
-            crate::chroma::send_to_chroma(client, &id, chunk, &embedding, &path, &pb).await?;
+            chunk_info.push((chunk.clone(), hash));
+        }
+
+        for batch in chunk_info.chunks(BATCH_SIZE) {
+            let mut tasks = FuturesUnordered::new();
+
+            for (chunk, hash) in batch.iter().cloned() {
+                let client = client.clone();
+                let path = path.to_path_buf();
+                let pb = pb.clone();
+                tasks.push(async move {
+                    sleep(Duration::from_millis(100)).await;
+                    let embedding = embedding::get_embedding(&client, &chunk).await?;
+                    let id = Uuid::new_v4().to_string();
+                    chroma::send_to_chroma(&client, &id, &chunk, &embedding, &path, &pb).await?;
+                    Ok::<_, anyhow::Error>(IndexedChunk { id, hash })
+                });
+            }
 
-            new_chunks.push(crate::state::IndexedChunk { id, hash });
+            while let Some(result) = tasks.next().await {
+                if let Ok(chunk) = result {
+                    new_chunks.push(chunk);
+                }
+            }
         }
 
         if !new_chunks.is_empty() {
             let mut updated_chunks = prev_chunks.clone();
-            let mut removed_chunks = vec![];
+            let mut removed_chunks = Vec::new();
 
             updated_chunks.retain(|c| {
                 let keep = new_chunks.iter().all(|n| n.hash != c.hash);
@@ -121,7 +146,7 @@ pub async fn handle_index(client: &Client, path: &Path) -> anyhow::Result<()> {
             state.update_file_chunks(&file_str, updated_chunks, modified);
 
             for chunk in removed_chunks {
-                crate::chroma::delete_chunk(client, &chunk.id).await?;
+                chroma::delete_chunk(client, &chunk.id).await?;
             }
         }
 

kb_core/src/cli/commands.rs

@@ -18,6 +18,47 @@ struct EmbeddingData {
     embedding: Vec<f32>,
 }
 
+pub async fn get_embeddings(client: &Client, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
+    let config = config::load_config()?;
+    let api_key = config::get_openai_api_key()?;
+
+    let body = serde_json::json!({
+        "model": config.openai_embedding_model,
+        "input": texts
+    });
+
+    let res = client
+        .post("https://api.openai.com/v1/embeddings")
+        .bearer_auth(api_key)
+        .json(&body)
+        .send()
+        .await?;
+
+    let status = res.status();
+    let text_body = res.text().await?;
+
+    if !status.is_success() {
+        eprintln!("❌ OpenAI error: HTTP {} - {}", status, text_body);
+        anyhow::bail!("Embedding batch failed");
+    }
+
+    let parsed: serde_json::Value = serde_json::from_str(&text_body)?;
+    let data = parsed["data"]
+        .as_array()
+        .ok_or_else(|| anyhow::anyhow!("Invalid embedding response format"))?;
+
+    Ok(data.iter()
+        .map(|v| {
+            v["embedding"]
+                .as_array()
+                .unwrap_or(&vec![])
+                .iter()
+                .map(|f| f.as_f64().unwrap_or_default() as f32)
+                .collect()
+        })
+        .collect())
+}
+
 pub async fn get_embedding(client: &Client, text: &str) -> anyhow::Result<Vec<f32>> {
     let config = config::load_config()?;
     let body = EmbeddingRequest {

kb_core/src/embedding/mod.rs

@@ -12,10 +12,8 @@ pub async fn get_llm_response(
     let config_dir = config::get_config_dir()?;
     let mut state = QueryState::load(&config_dir)?;
 
-    // Generate context hash from context chunks
     let context_hash = hash_query_context(prompt, context_chunks);
 
-    // Return cached answer if it exists
     if let Some(cached) = state.get_cached_answer(prompt, &context_hash) {
         return Ok(cached);
     }

kb_core/src/llm/mod.rs

@@ -58,7 +58,7 @@ pub fn collect_files(root: &Path) -> anyhow::Result<Vec<PathBuf>> {
         }
     } else {
         let walker = WalkBuilder::new(root)
-            .add_custom_ignore_filename(".kbignore") // Optional
+            .add_custom_ignore_filename(".kbignore")
             .hidden(false)
             .build();