added caching

Author: thesandybridge

Cargo.lock

@@ -25,3 +25,5 @@ syntect-assets = "0.23.6"
 lazy_static = "1.5.0"
 bat = { version = "0.25.0", features = ["build-assets"] }
 ignore = "0.4.23"
+sha2 = "0.10.9"
+hex = "0.4.3"

kb_core/Cargo.toml

@@ -159,3 +159,27 @@ pub async fn query_chroma(
 
     Ok(parsed)
 }
+
+pub async fn delete_chunk(client: &Client, id: &str) -> anyhow::Result<()> {
+    let config = config::load_config()?;
+    let collection_id = get_collection_id(client).await?;
+
+    let url = format!(
+        "{}/api/v2/tenants/{}/databases/{}/collections/{}/delete",
+        config.chroma_host, TENANT, DATABASE, collection_id
+    );
+
+    let payload = serde_json::json!({
+        "ids": [id]
+    });
+
+    let resp = client.post(&url).json(&payload).send().await?;
+    let status = resp.status();
+    let body = resp.text().await?;
+
+    if !status.is_success() {
+        anyhow::bail!("Failed to delete chunk {}: HTTP {} - {}", id, status, body);
+    }
+
+    Ok(())
+}

kb_core/src/chroma/mod.rs

@@ -3,6 +3,7 @@ use crate::config;
 use crate::embedding;
 use crate::llm;
 use crate::utils;
+use std::time::UNIX_EPOCH;
 use indicatif::{ProgressBar, ProgressStyle};
 use reqwest::Client;
 use std::fs;
@@ -56,7 +57,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 = utils::collect_files(path)?;
+    let paths = crate::utils::collect_files(path)?;
     let total_files = paths.len() as u64;
     let pb = ProgressBar::new(total_files);
     pb.set_style(
@@ -65,23 +66,69 @@ 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)?;
+
     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
+        if let Some(prev) = state.get_last_modified(&file_str) {
+            if prev == modified {
+                pb.inc(1);
+                continue;
+            }
+        }
+
         let content = fs::read_to_string(&path)?;
-        let chunks = utils::chunk_text(&content);
+        let chunks = crate::utils::chunk_text(&content);
+        let prev_chunks = state.get_file_chunks(&file_str).cloned().unwrap_or_default();
+        let mut new_chunks = vec![];
 
         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) {
+                continue;
+            }
+
             let id = Uuid::new_v4().to_string();
-            let embedding = embedding::get_embedding(&client, &chunk).await?;
-            chroma::send_to_chroma(&client, &id, &chunk, &embedding, &path, &pb).await?;
+            let embedding = crate::embedding::get_embedding(client, chunk).await?;
+            crate::chroma::send_to_chroma(client, &id, chunk, &embedding, &path, &pb).await?;
+
+            new_chunks.push(crate::state::IndexedChunk { id, hash });
         }
+
+        if !new_chunks.is_empty() {
+            let mut updated_chunks = prev_chunks.clone();
+            let mut removed_chunks = vec![];
+
+            updated_chunks.retain(|c| {
+                let keep = new_chunks.iter().all(|n| n.hash != c.hash);
+                if !keep {
+                    removed_chunks.push(c.clone());
+                }
+                keep
+            });
+
+            updated_chunks.extend(new_chunks);
+            state.update_file_chunks(&file_str, updated_chunks, modified);
+
+            for chunk in removed_chunks {
+                crate::chroma::delete_chunk(client, &chunk.id).await?;
+            }
+        }
+
         pb.inc(1);
     }
 
+    state.save(&config_dir)?;
     pb.finish_with_message("🎉 Indexing complete.");
     Ok(())
 }
@@ -149,7 +196,7 @@ pub async fn handle_query(
             }
         }
         "smart" => {
-            let annotated_chunks: Vec<String> = results.iter()
+            let context_chunks: Vec<String> = results.iter()
                 .map(|r| {
                     let lang = Path::new(r.source)
                         .extension()
@@ -163,20 +210,9 @@ pub async fn handle_query(
                 })
                 .collect();
 
-            let context = annotated_chunks.join("\n\n---\n\n");
-
-            let prompt = format!(
-                "You are a helpful coding and personal assistant.\n\
-                    Use the following code snippets to answer the question. Each file is shown with its source path and syntax.\n\
-                    Format your response in Markdown and include code where necessary.\n\n\
-                    Question:\n{}\n\n\
-                    Documents:\n{}\n\n\
-                    Answer:",
-                query, context
-            );
-
-            let raw_answer = llm::get_llm_response(client, &prompt).await?;
+            let raw_answer = llm::get_llm_response(client, query, &context_chunks).await?;
             let rendered = utils::render_markdown_highlighted(&raw_answer);
+
             println!("💡 Answer:\n\n{}", rendered);
         }
         _ => {

kb_core/src/cli/commands.rs

@@ -2,6 +2,7 @@ use dirs::config_dir;
 use serde::{Deserialize, Serialize};
 use std::env;
 use std::fs;
+use std::path::PathBuf;
 
 #[derive(Deserialize, Serialize, Clone)]
 pub struct AppConfig {
@@ -81,3 +82,12 @@ pub fn get_openai_api_key() -> anyhow::Result<String> {
 
     anyhow::bail!("OpenAI API key not found in environment or config file. Please set the OPENAI_API_KEY environment variable or add it to your config file.")
 }
+
+pub fn get_config_dir() -> anyhow::Result<PathBuf> {
+    let path = config_dir()
+        .ok_or_else(|| anyhow::anyhow!("Unable to determine config directory"))?
+        .join("kb-index");
+
+    std::fs::create_dir_all(&path)?;
+    Ok(path)
+}

kb_core/src/config/mod.rs

@@ -4,3 +4,4 @@ pub mod config;
 pub mod embedding;
 pub mod utils;
 pub mod llm;
+pub mod state;

kb_core/src/lib.rs

@@ -1,15 +1,41 @@
-use reqwest::Client;
 use crate::config;
+use crate::state::{QueryState, hash_query_context};
+use reqwest::Client;
 
-pub async fn get_llm_response(client: &Client, prompt: &str) -> anyhow::Result<String> {
+pub async fn get_llm_response(
+    client: &Client,
+    prompt: &str,
+    context_chunks: &[String],
+) -> anyhow::Result<String> {
     let api_key = config::get_openai_api_key()?;
-    let config = config::load_config()?;
+    let cfg = config::load_config()?;
+    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);
+    }
+
+    // Concatenate all context chunks for the prompt
+    let full_context = context_chunks.join("\n\n---\n\n");
+
+    let full_prompt = format!(
+        "You are an expert personal and code assistant.\n\
+         Use the following code snippets to answer the question. \
+         Format your response in Markdown and include code where necessary.\n\n\
+         Question:\n{}\n\nContext:\n{}\n\nAnswer:",
+        prompt, full_context
+    );
 
     let body = serde_json::json!({
-        "model": config.openai_completion_model,
+        "model": cfg.openai_completion_model,
         "messages": [
-            { "role": "system", "content": "You are a expert personal and code assistant." },
-            { "role": "user", "content": prompt }
+            { "role": "system", "content": "You are an expert personal and code assistant." },
+            { "role": "user", "content": full_prompt }
         ],
         "temperature": 0.4
     });
@@ -29,5 +55,9 @@ pub async fn get_llm_response(client: &Client, prompt: &str) -> anyhow::Result<S
         .unwrap_or("No answer generated")
         .to_string();
 
+    // Cache the answer
+    state.insert_answer(prompt.to_string(), context_hash, answer.clone());
+    state.save(&config_dir)?;
+
     Ok(answer)
 }