370 lines · 14.1 KB · Rust
RawDownload
use anyhow::{anyhow, Result};use assistant_slash_command::{    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,    SlashCommandResult,};use feature_flags::FeatureFlag;use futures::StreamExt;use gpui::{AppContext, AsyncAppContext, AsyncWindowContext, Task, WeakView, WindowContext};use language::{CodeLabel, LspAdapterDelegate};use language_model::{    LanguageModelCompletionEvent, LanguageModelRegistry, LanguageModelRequest,    LanguageModelRequestMessage, Role,};use semantic_index::{FileSummary, SemanticDb};use smol::channel;use std::sync::{atomic::AtomicBool, Arc};use ui::{prelude::*, BorrowAppContext};use util::ResultExt;use workspace::Workspace; use crate::slash_command::create_label_for_command; pub struct AutoSlashCommandFeatureFlag; impl FeatureFlag for AutoSlashCommandFeatureFlag {    const NAME: &'static str = "auto-slash-command";} pub(crate) struct AutoCommand; impl SlashCommand for AutoCommand {    fn name(&self) -> String {        "auto".into()    }     fn description(&self) -> String {        "Automatically infer what context to add".into()    }     fn icon(&self) -> IconName {        IconName::Wand    }     fn menu_text(&self) -> String {        self.description()    }     fn label(&self, cx: &AppContext) -> CodeLabel {        create_label_for_command("auto", &["--prompt"], cx)    }     fn complete_argument(        self: Arc<Self>,        _arguments: &[String],        _cancel: Arc<AtomicBool>,        workspace: Option<WeakView<Workspace>>,        cx: &mut WindowContext,    ) -> Task<Result<Vec<ArgumentCompletion>>> {        // There's no autocomplete for a prompt, since it's arbitrary text.        // However, we can use this opportunity to kick off a drain of the backlog.        // That way, it can hopefully be done resummarizing by the time we've actually        // typed out our prompt. This re-runs on every keystroke during autocomplete,        // but in the future, we could instead do it only once, when /auto is first entered.        let Some(workspace) = workspace.and_then(|ws| ws.upgrade()) else {            log::warn!("workspace was dropped or unavailable during /auto autocomplete");             return Task::ready(Ok(Vec::new()));        };         let project = workspace.read(cx).project().clone();        let Some(project_index) =            cx.update_global(|index: &mut SemanticDb, cx| index.project_index(project, cx))        else {            return Task::ready(Err(anyhow!("No project indexer, cannot use /auto")));        };         let cx: &mut AppContext = cx;         cx.spawn(|cx: gpui::AsyncAppContext| async move {            let task = project_index.read_with(&cx, |project_index, cx| {                project_index.flush_summary_backlogs(cx)            })?;             cx.background_executor().spawn(task).await;             anyhow::Ok(Vec::new())        })    }     fn requires_argument(&self) -> bool {        true    }     fn run(        self: Arc<Self>,        arguments: &[String],        _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],        _context_buffer: language::BufferSnapshot,        workspace: WeakView<Workspace>,        _delegate: Option<Arc<dyn LspAdapterDelegate>>,        cx: &mut WindowContext,    ) -> Task<SlashCommandResult> {        let Some(workspace) = workspace.upgrade() else {            return Task::ready(Err(anyhow::anyhow!("workspace was dropped")));        };        if arguments.is_empty() {            return Task::ready(Err(anyhow!("missing prompt")));        };        let argument = arguments.join(" ");        let original_prompt = argument.to_string();        let project = workspace.read(cx).project().clone();        let Some(project_index) =            cx.update_global(|index: &mut SemanticDb, cx| index.project_index(project, cx))        else {            return Task::ready(Err(anyhow!("no project indexer")));        };         let task = cx.spawn(|cx: AsyncWindowContext| async move {            let summaries = project_index                .read_with(&cx, |project_index, cx| project_index.all_summaries(cx))?                .await?;             commands_for_summaries(&summaries, &original_prompt, &cx).await        });         // As a convenience, append /auto's argument to the end of the prompt        // so you don't have to write it again.        let original_prompt = argument.to_string();         cx.background_executor().spawn(async move {            let commands = task.await?;            let mut prompt = String::new();             log::info!(                "Translating this response into slash-commands: {:?}",                commands            );             for command in commands {                prompt.push('/');                prompt.push_str(&command.name);                prompt.push(' ');                prompt.push_str(&command.arg);                prompt.push('\n');            }             prompt.push('\n');            prompt.push_str(&original_prompt);             Ok(SlashCommandOutput {                text: prompt,                sections: Vec::new(),                run_commands_in_text: true,            }            .to_event_stream())        })    }} const PROMPT_INSTRUCTIONS_BEFORE_SUMMARY: &str = include_str!("prompt_before_summary.txt");const PROMPT_INSTRUCTIONS_AFTER_SUMMARY: &str = include_str!("prompt_after_summary.txt"); fn summaries_prompt(summaries: &[FileSummary], original_prompt: &str) -> String {    let json_summaries = serde_json::to_string(summaries).unwrap();     format!("{PROMPT_INSTRUCTIONS_BEFORE_SUMMARY}\n{json_summaries}\n{PROMPT_INSTRUCTIONS_AFTER_SUMMARY}\n{original_prompt}")} /// The slash commands that the model is told about, and which we look for in the inference response.const SUPPORTED_SLASH_COMMANDS: &[&str] = &["search", "file"]; #[derive(Debug, Clone)]struct CommandToRun {    name: String,    arg: String,} /// Given the pre-indexed file summaries for this project, as well as the original prompt/// string passed to `/auto`, get a list of slash commands to run, along with their arguments.////// The prompt's output does not include the slashes (to reduce the chance that it makes a mistake),/// so taking one of these returned Strings and turning it into a real slash-command-with-argument/// involves prepending a slash to it.////// This function will validate that each of the returned lines begins with one of SUPPORTED_SLASH_COMMANDS./// Any other lines it encounters will be discarded, with a warning logged.async fn commands_for_summaries(    summaries: &[FileSummary],    original_prompt: &str,    cx: &AsyncAppContext,) -> Result<Vec<CommandToRun>> {    if summaries.is_empty() {        log::warn!("Inferring no context because there were no summaries available.");        return Ok(Vec::new());    }     // Use the globally configured model to translate the summaries into slash-commands,    // because Qwen2-7B-Instruct has not done a good job at that task.    let Some(model) = cx.update(|cx| LanguageModelRegistry::read_global(cx).active_model())? else {        log::warn!("Can't infer context because there's no active model.");        return Ok(Vec::new());    };    // Only go up to 90% of the actual max token count, to reduce chances of    // exceeding the token count due to inaccuracies in the token counting heuristic.    let max_token_count = (model.max_token_count() * 9) / 10;     // Rather than recursing (which would require this async function use a pinned box),    // we use an explicit stack of arguments and answers for when we need to "recurse."    let mut stack = vec![summaries];    let mut final_response = Vec::new();    let mut prompts = Vec::new();     // TODO We only need to create multiple Requests because we currently    // don't have the ability to tell if a CompletionProvider::complete response    // was a "too many tokens in this request" error. If we had that, then    // we could try the request once, instead of having to make separate requests    // to check the token count and then afterwards to run the actual prompt.    let make_request = |prompt: String| LanguageModelRequest {        messages: vec![LanguageModelRequestMessage {            role: Role::User,            content: vec![prompt.into()],            // Nothing in here will benefit from caching            cache: false,        }],        tools: Vec::new(),        stop: Vec::new(),        temperature: None,    };     while let Some(current_summaries) = stack.pop() {        // The split can result in one slice being empty and the other having one element.        // Whenever that happens, skip the empty one.        if current_summaries.is_empty() {            continue;        }         log::info!(            "Inferring prompt context using {} file summaries",            current_summaries.len()        );         let prompt = summaries_prompt(&current_summaries, original_prompt);        let start = std::time::Instant::now();        // Per OpenAI, 1 token ~= 4 chars in English (we go with 4.5 to overestimate a bit, because failed API requests cost a lot of perf)        // Verifying this against an actual model.count_tokens() confirms that it's usually within ~5% of the correct answer, whereas        // getting the correct answer from tiktoken takes hundreds of milliseconds (compared to this arithmetic being ~free).        // source: https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them        let token_estimate = prompt.len() * 2 / 9;        let duration = start.elapsed();        log::info!(            "Time taken to count tokens for prompt of length {:?}B: {:?}",            prompt.len(),            duration        );         if token_estimate < max_token_count {            prompts.push(prompt);        } else if current_summaries.len() == 1 {            log::warn!("Inferring context for a single file's summary failed because the prompt's token length exceeded the model's token limit.");        } else {            log::info!(                "Context inference using file summaries resulted in a prompt containing {token_estimate} tokens, which exceeded the model's max of {max_token_count}. Retrying as two separate prompts, each including half the number of summaries.",            );            let (left, right) = current_summaries.split_at(current_summaries.len() / 2);            stack.push(right);            stack.push(left);        }    }     let all_start = std::time::Instant::now();     let (tx, rx) = channel::bounded(1024);     let completion_streams = prompts        .into_iter()        .map(|prompt| {            let request = make_request(prompt.clone());            let model = model.clone();            let tx = tx.clone();            let stream = model.stream_completion(request, &cx);             (stream, tx)        })        .collect::<Vec<_>>();     cx.background_executor()        .spawn(async move {            let futures = completion_streams                .into_iter()                .enumerate()                .map(|(ix, (stream, tx))| async move {                    let start = std::time::Instant::now();                    let events = stream.await?;                    log::info!("Time taken for awaiting /await chunk stream #{ix}: {:?}", start.elapsed());                     let completion: String = events                        .filter_map(|event| async {                            if let Ok(LanguageModelCompletionEvent::Text(text)) = event {                                Some(text)                            } else {                                None                            }                        })                        .collect()                        .await;                     log::info!("Time taken for all /auto chunks to come back for #{ix}: {:?}", start.elapsed());                     for line in completion.split('\n') {                        if let Some(first_space) = line.find(' ') {                            let command = &line[..first_space].trim();                            let arg = &line[first_space..].trim();                             tx.send(CommandToRun {                                name: command.to_string(),                                arg: arg.to_string(),                            })                            .await?;                        } else if !line.trim().is_empty() {                            // All slash-commands currently supported in context inference need a space for the argument.                            log::warn!(                                "Context inference returned a non-blank line that contained no spaces (meaning no argument for the slash command): {:?}",                                line                            );                        }                    }                     anyhow::Ok(())                })                .collect::<Vec<_>>();             let _ = futures::future::try_join_all(futures).await.log_err();             let duration = all_start.elapsed();            eprintln!("All futures completed in {:?}", duration);        })        .await;     drop(tx); // Close the channel so that rx.collect() won't hang. This is safe because all futures have completed.    let results = rx.collect::<Vec<_>>().await;    eprintln!(        "Finished collecting from the channel with {} results",        results.len()    );    for command in results {        // Don't return empty or duplicate commands        if !command.name.is_empty()            && !final_response                .iter()                .any(|cmd: &CommandToRun| cmd.name == command.name && cmd.arg == command.arg)        {            if SUPPORTED_SLASH_COMMANDS                .iter()                .any(|supported| &command.name == supported)            {                final_response.push(command);            } else {                log::warn!(                    "Context inference returned an unrecognized slash command: {:?}",                    command                );            }        }    }     // Sort the commands by name (reversed just so that /search appears before /file)    final_response.sort_by(|cmd1, cmd2| cmd1.name.cmp(&cmd2.name).reverse());     Ok(final_response)}