use crate::{assistant_settings::OpenAiModel, MessageId, MessageMetadata}; use anyhow::{anyhow, Result}; use assistant_slash_command::SlashCommandOutputSection; use collections::HashMap; use fs::Fs; use futures::StreamExt; use fuzzy::StringMatchCandidate; use gpui::{AppContext, Model, ModelContext, Task}; use paths::CONTEXTS_DIR; use regex::Regex; use serde::{Deserialize, Serialize}; use std::{cmp::Reverse, ffi::OsStr, path::PathBuf, sync::Arc, time::Duration}; use ui::Context; use util::{ResultExt, TryFutureExt}; #[derive(Serialize, Deserialize)] pub struct SavedMessage { pub id: MessageId, pub start: usize, } #[derive(Serialize, Deserialize)] pub struct SavedContext { pub id: Option, pub zed: String, pub version: String, pub text: String, pub messages: Vec, pub message_metadata: HashMap, pub summary: String, pub slash_command_output_sections: Vec>, } impl SavedContext { pub const VERSION: &'static str = "0.3.0"; } #[derive(Serialize, Deserialize)] pub struct SavedContextV0_2_0 { pub id: Option, pub zed: String, pub version: String, pub text: String, pub messages: Vec, pub message_metadata: HashMap, pub summary: String, } #[derive(Serialize, Deserialize)] struct SavedContextV0_1_0 { id: Option, zed: String, version: String, text: String, messages: Vec, message_metadata: HashMap, summary: String, api_url: Option, model: OpenAiModel, } #[derive(Clone)] pub struct SavedContextMetadata { pub title: String, pub path: PathBuf, pub mtime: chrono::DateTime, } pub struct ContextStore { contexts_metadata: Vec, fs: Arc, _watch_updates: Task>, } impl ContextStore { pub fn new(fs: Arc, cx: &mut AppContext) -> Task>> { cx.spawn(|mut cx| async move { const CONTEXT_WATCH_DURATION: Duration = Duration::from_millis(100); let (mut events, _) = fs.watch(&CONTEXTS_DIR, CONTEXT_WATCH_DURATION).await; let this = cx.new_model(|cx: &mut ModelContext| Self { contexts_metadata: Vec::new(), fs, _watch_updates: cx.spawn(|this, mut cx| { async move { while events.next().await.is_some() { this.update(&mut cx, |this, cx| this.reload(cx))? .await .log_err(); } anyhow::Ok(()) } .log_err() }), })?; this.update(&mut cx, |this, cx| this.reload(cx))? .await .log_err(); Ok(this) }) } pub fn load(&self, path: PathBuf, cx: &AppContext) -> Task> { let fs = self.fs.clone(); cx.background_executor().spawn(async move { let saved_context = fs.load(&path).await?; let saved_context_json = serde_json::from_str::(&saved_context)?; match saved_context_json .get("version") .ok_or_else(|| anyhow!("version not found"))? { serde_json::Value::String(version) => match version.as_str() { SavedContext::VERSION => { Ok(serde_json::from_value::(saved_context_json)?) } "0.2.0" => { let saved_context = serde_json::from_value::(saved_context_json)?; Ok(SavedContext { id: saved_context.id, zed: saved_context.zed, version: saved_context.version, text: saved_context.text, messages: saved_context.messages, message_metadata: saved_context.message_metadata, summary: saved_context.summary, slash_command_output_sections: Vec::new(), }) } "0.1.0" => { let saved_context = serde_json::from_value::(saved_context_json)?; Ok(SavedContext { id: saved_context.id, zed: saved_context.zed, version: saved_context.version, text: saved_context.text, messages: saved_context.messages, message_metadata: saved_context.message_metadata, summary: saved_context.summary, slash_command_output_sections: Vec::new(), }) } _ => Err(anyhow!("unrecognized saved context version: {}", version)), }, _ => Err(anyhow!("version not found on saved context")), } }) } pub fn search(&self, query: String, cx: &AppContext) -> Task> { let metadata = self.contexts_metadata.clone(); let executor = cx.background_executor().clone(); cx.background_executor().spawn(async move { if query.is_empty() { metadata } else { let candidates = metadata .iter() .enumerate() .map(|(id, metadata)| StringMatchCandidate::new(id, metadata.title.clone())) .collect::>(); let matches = fuzzy::match_strings( &candidates, &query, false, 100, &Default::default(), executor, ) .await; matches .into_iter() .map(|mat| metadata[mat.candidate_id].clone()) .collect() } }) } fn reload(&mut self, cx: &mut ModelContext) -> Task> { let fs = self.fs.clone(); cx.spawn(|this, mut cx| async move { fs.create_dir(&CONTEXTS_DIR).await?; let mut paths = fs.read_dir(&CONTEXTS_DIR).await?; let mut contexts = Vec::::new(); while let Some(path) = paths.next().await { let path = path?; if path.extension() != Some(OsStr::new("json")) { continue; } let pattern = r" - \d+.zed.json$"; let re = Regex::new(pattern).unwrap(); let metadata = fs.metadata(&path).await?; if let Some((file_name, metadata)) = path .file_name() .and_then(|name| name.to_str()) .zip(metadata) { // This is used to filter out contexts saved by the new assistant. if !re.is_match(file_name) { continue; } if let Some(title) = re.replace(file_name, "").lines().next() { contexts.push(SavedContextMetadata { title: title.to_string(), path, mtime: metadata.mtime.into(), }); } } } contexts.sort_unstable_by_key(|context| Reverse(context.mtime)); this.update(&mut cx, |this, cx| { this.contexts_metadata = contexts; cx.notify(); }) }) } }