added support for custom sessions

Author: thesandybridge

Cargo.lock

@@ -28,3 +28,4 @@ ignore = "0.4.23"
 sha2 = "0.10.9"
 hex = "0.4.3"
 futures = "0.3.31"
+chrono = "0.4.41"

kb_core/Cargo.toml

@@ -1,3 +1,4 @@
 pub mod config;
 pub mod index;
 pub mod query;
+pub mod session;

kb_core/src/cli/commands/mod.rs

@@ -4,24 +4,51 @@ use crate::llm;
 use crate::utils;
 use reqwest::Client;
 use std::path::Path;
-use crate::state::{QueryState, hash_query_context};
+use crate::state::{QueryState, SessionManager, hash_query_context};
 use crate::config;
 
 pub async fn handle_query(
     client: &Client,
     query: &str,
     top_k: usize,
     format: &str,
+    session_id: Option<String>,
 ) -> anyhow::Result<()> {
     let config_dir = config::get_config_dir()?;
     let mut cache = QueryState::load(&config_dir)?;
+    let mut session_manager = SessionManager::load(&config_dir)?;
+
+    // Handle session management
+    if let Some(id) = session_id {
+        if id == "new" {
+            let new_id = session_manager.create_session();
+            println!("🆕 Created new session: {}", new_id);
+        } else {
+            // Pass a reference to set_active_session
+            session_manager.set_active_session(&id)?;
+            println!("🔄 Switched to session: {}", id);
+        }
+    } else if session_manager.active_session.is_none() {
+        // Create a default session if none exists
+        let new_id = session_manager.create_session();
+        println!("🆕 Created default session: {}", new_id);
+    }
+
 
     // Embed the query
     let query_embedding = embedding::get_embedding(client, query).await?;
 
     // 🔍 Try similarity cache
     if let Some(similar) = cache.find_similar(&query_embedding, 0.93) {
         println!("💡 Cached Answer:\n\n{}", utils::render_markdown_highlighted(&similar));
+
+        // Add to session history even if cached
+        if let Some(session) = session_manager.get_active_session_mut() {
+            session.queries.push(query.to_string());
+            session.responses.push(similar);
+            session_manager.save(&config_dir)?;
+        }
+
         return Ok(());
     }
 
@@ -97,14 +124,33 @@ pub async fn handle_query(
                 .collect();
 
             let context_hash = hash_query_context(query, &context_chunks);
-            let raw_answer = llm::get_llm_response(client, query, &context_chunks).await?;
+
+            // Pass session manager to get_llm_response
+            let raw_answer = llm::get_llm_response(
+                client,
+                query,
+                &context_chunks,
+                Some(&session_manager)
+            ).await?;
+
             let rendered = utils::render_markdown_highlighted(&raw_answer);
 
             // 🧠 Cache the answer with the current query embedding
             cache.insert_answer(query.to_string(), context_hash, query_embedding.clone(), raw_answer.clone());
             cache.save(&config_dir)?;
 
+            // Add to session history
+            session_manager.add_interaction(query.to_string(), raw_answer)?;
+            session_manager.save(&config_dir)?;
+
             println!("💡 Answer:\n\n{}", rendered);
+
+            if let Some(session) = session_manager.get_active_session() {
+                println!("\n📝 Session: {} (Q&A: {})",
+                    &session.id[..8],
+                    session.queries.len()
+                );
+            }
         }
         _ => {
             for r in &results {
@@ -119,3 +165,4 @@ pub async fn handle_query(
 
     Ok(())
 }
+

kb_core/src/cli/commands/query.rs

@@ -0,0 +1,54 @@
+use crate::{config, state::SessionManager};
+
+
+pub fn handle_sessions(list: bool, clear: bool, switch: Option<String>) -> anyhow::Result<()> {
+    let config_dir = config::get_config_dir()?;
+    let mut session_manager = SessionManager::load(&config_dir)?;
+
+    if clear {
+        if let Some(active_id) = session_manager.active_session.clone() {
+            session_manager.sessions.remove(&active_id);
+            session_manager.active_session = None;
+            session_manager.save(&config_dir)?;
+            println!("🧹 Cleared session: {}", active_id);
+        } else {
+            println!("⚠️ No active session to clear");
+        }
+        return Ok(());
+    }
+
+    if let Some(id) = switch {
+        session_manager.set_active_session(&id)?;
+        println!("🔄 Switched to session: {}", id);
+        session_manager.save(&config_dir)?;
+        return Ok(());
+    }
+
+    if list || session_manager.sessions.is_empty() {
+        println!("📋 Available Sessions:");
+        for (id, session) in session_manager.list_sessions() {
+            let active = if Some(id) == session_manager.active_session.as_ref() {
+                "* "
+            } else {
+                "  "
+            };
+
+            let time = chrono::DateTime::<chrono::Utc>::from_timestamp(
+                session.last_updated as i64, 0
+            ).unwrap_or_default().format("%Y-%m-%d %H:%M");
+
+            println!("{}{} - {} Q&A pairs, last updated: {}",
+                active,
+                &id[..8],
+                session.queries.len(),
+                time
+            );
+        }
+
+        if session_manager.sessions.is_empty() {
+            println!("  No sessions found. Create one with 'kb query --session new'");
+        }
+    }
+
+    Ok(())
+}

kb_core/src/cli/commands/session.rs

@@ -19,9 +19,24 @@ pub enum Cli {
         /// Number of results to return
         #[arg(short, long, default_value_t = 5)]
         top_k: usize,
-        /// Output format (pretty, json, markdown)
+        /// Output format (pretty, json, markdown, smart)
         #[arg(short, long, default_value = "smart")]
         format: String,
+        /// Session ID or 'new' to create a new session
+        #[arg(long)]
+        session: Option<String>,
+    },
+    /// Manage sessions for conversation history
+    Sessions {
+        /// List all available sessions
+        #[arg(short, long, default_value_t = false)]
+        list: bool,
+        /// Clear the current session
+        #[arg(short, long, default_value_t = false)]
+        clear: bool,
+        /// Switch to a specific session by ID
+        #[arg(short, long)]
+        switch: Option<String>,
     },
     /// Configure the application
     Config {
@@ -33,3 +48,4 @@ pub enum Cli {
         show: bool,
     },
 }
+

kb_core/src/cli/mod.rs

@@ -1,12 +1,14 @@
 use crate::config;
 use crate::embedding;
+use crate::state::SessionManager;
 use crate::state::{QueryState, hash_query_context};
 use reqwest::Client;
 
 pub async fn get_llm_response(
     client: &Client,
     prompt: &str,
     context_chunks: &[String],
+    session_manager: Option<&SessionManager>,
 ) -> anyhow::Result<String> {
     let api_key = config::get_openai_api_key()?;
     let cfg = config::load_config()?;
@@ -21,27 +23,56 @@ pub async fn get_llm_response(
     if let Some(similar) = state.find_similar(&embedding, 0.93) {
         return Ok(similar);
     }
+
     // Check for cached similar answer
     if let Some(cached) = state.get_cached_answer(prompt, &context_hash) {
         return Ok(cached);
     }
 
-    // Prepare full prompt
+    // Prepare full prompt with session context if available
     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. \
+
+    let mut messages = vec![
+        serde_json::json!({
+            "role": "system",
+            "content": "You are an expert personal and code assistant."
+        }),
+    ];
+
+    // Add session context if available
+    if let Some(manager) = session_manager {
+        if let Some(session) = manager.get_active_session() {
+            // Add previous interactions as context
+            for (q, r) in session.queries.iter().zip(session.responses.iter()) {
+                messages.push(serde_json::json!({
+                    "role": "user",
+                    "content": q
+                }));
+
+                messages.push(serde_json::json!({
+                    "role": "assistant",
+                    "content": r
+                }));
+            }
+        }
+    }
+
+    // Add current query with context
+    let user_content = format!(
+        "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:",
+         Question:\n{}\n\nContext:\n{}",
         prompt, full_context
     );
 
+    messages.push(serde_json::json!({
+        "role": "user",
+        "content": user_content
+    }));
+
     let body = serde_json::json!({
         "model": cfg.openai_completion_model,
-        "messages": [
-            { "role": "system", "content": "You are an expert personal and code assistant." },
-            { "role": "user", "content": full_prompt }
-        ],
+        "messages": messages,
         "temperature": 0.4
     });
 
@@ -65,3 +96,4 @@ pub async fn get_llm_response(
 
     Ok(answer)
 }
+