From ad3055076d146c49c1986286c142b4bee5bad99c Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Thu, 18 Jul 2024 17:01:48 -0400 Subject: [PATCH] assistant: Allow `/docs` to perform JIT indexing when run (#14768) This PR updates the `/docs` slash command with the ability to just-in-time index a package when there are not yet any results in the index. When running a `/docs` slash command, we fist check to see if there are any results in the index that would match the search. If there are, we go ahead and return them, as we do today. However, if there are not yet any results we kick off an indexing task as part of the command execution to fetch the results. Release Notes: - N/A --- .../src/slash_command/docs_command.rs | 63 ++++++++++++++++++- crates/indexed_docs/src/store.rs | 24 +++++++ 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/crates/assistant/src/slash_command/docs_command.rs b/crates/assistant/src/slash_command/docs_command.rs index 9fb57a0d9591cf308ae49c11e6f632db579329ec..6271bdc32e6b412b3fd66135ab2f384961894727 100644 --- a/crates/assistant/src/slash_command/docs_command.rs +++ b/crates/assistant/src/slash_command/docs_command.rs @@ -1,12 +1,13 @@ use std::path::Path; use std::sync::atomic::AtomicBool; use std::sync::Arc; +use std::time::Duration; use anyhow::{anyhow, bail, Result}; use assistant_slash_command::{ ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection, }; -use gpui::{AppContext, Model, Task, WeakView}; +use gpui::{AppContext, BackgroundExecutor, Model, Task, WeakView}; use indexed_docs::{ DocsDotRsProvider, IndexedDocsRegistry, IndexedDocsStore, LocalRustdocProvider, PackageName, ProviderId, @@ -90,6 +91,55 @@ impl DocsSlashCommand { } } } + + /// Runs just-in-time indexing for a given package, in case the slash command + /// is run without any entries existing in the index. + fn run_just_in_time_indexing( + store: Arc, + key: String, + package: PackageName, + executor: BackgroundExecutor, + ) -> Task<()> { + executor.clone().spawn(async move { + let (prefix, needs_full_index) = if let Some((prefix, _)) = key.split_once('*') { + // If we have a wildcard in the search, we want to wait until + // we've completely finished indexing so we get a full set of + // results for the wildcard. + (prefix.to_string(), true) + } else { + (key, false) + }; + + // If we already have some entries, we assume that we've indexed the package before + // and don't need to do it again. + let has_any_entries = store + .any_with_prefix(prefix.clone()) + .await + .unwrap_or_default(); + if has_any_entries { + return (); + }; + + let index_task = store.clone().index(package.clone()); + + if needs_full_index { + _ = index_task.await; + } else { + loop { + executor.timer(Duration::from_millis(200)).await; + + if store + .any_with_prefix(prefix.clone()) + .await + .unwrap_or_default() + || !store.is_indexing(&package) + { + break; + } + } + } + }) + } } impl SlashCommand for DocsSlashCommand { @@ -200,13 +250,14 @@ impl SlashCommand for DocsSlashCommand { }; let args = DocsSlashCommandArgs::parse(argument); + let executor = cx.background_executor().clone(); let task = cx.background_executor().spawn({ let store = args .provider() .ok_or_else(|| anyhow!("no docs provider specified")) .and_then(|provider| IndexedDocsStore::try_global(provider, cx)); async move { - let (provider, key) = match args { + let (provider, key) = match args.clone() { DocsSlashCommandArgs::NoProvider => bail!("no docs provider specified"), DocsSlashCommandArgs::SearchPackageDocs { provider, package, .. @@ -219,6 +270,12 @@ impl SlashCommand for DocsSlashCommand { }; let store = store?; + + if let Some(package) = args.package() { + Self::run_just_in_time_indexing(store.clone(), key.clone(), package, executor) + .await; + } + let (text, ranges) = if let Some((prefix, _)) = key.split_once('*') { let docs = store.load_many_by_prefix(prefix.to_string()).await?; @@ -269,7 +326,7 @@ fn is_item_path_delimiter(char: char) -> bool { !char.is_alphanumeric() && char != '-' && char != '_' } -#[derive(Debug, PartialEq)] +#[derive(Debug, PartialEq, Clone)] pub(crate) enum DocsSlashCommandArgs { NoProvider, SearchPackageDocs { diff --git a/crates/indexed_docs/src/store.rs b/crates/indexed_docs/src/store.rs index 4f5ee10f057f432d761cf68ea5f9d234fd9444ad..b581e93ec0a023277aa4575949a61c10776777ba 100644 --- a/crates/indexed_docs/src/store.rs +++ b/crates/indexed_docs/src/store.rs @@ -112,6 +112,16 @@ impl IndexedDocsStore { .await } + /// Returns whether any entries exist with the given prefix. + pub async fn any_with_prefix(&self, prefix: String) -> Result { + self.database_future + .clone() + .await + .map_err(|err| anyhow!(err))? + .any_with_prefix(prefix) + .await + } + pub fn index( self: Arc, package: PackageName, @@ -288,6 +298,20 @@ impl IndexedDocsDatabase { }) } + /// Returns whether any entries exist with the given prefix. + pub fn any_with_prefix(&self, prefix: String) -> Task> { + let env = self.env.clone(); + let entries = self.entries; + + self.executor.spawn(async move { + let txn = env.read_txn()?; + let any = entries + .iter(&txn)? + .any(|entry| entry.map_or(false, |(key, _value)| key.starts_with(&prefix))); + Ok(any) + }) + } + pub fn insert(&self, key: String, docs: String) -> Task> { let env = self.env.clone(); let entries = self.entries;