304 lines · 9.8 KB · Rust
RawDownload
use anyhow::Result;use assistant_tooling::{LanguageModelTool, ProjectContext, ToolView};use editor::{    display_map::{BlockContext, BlockDisposition, BlockProperties, BlockStyle},    Editor, MultiBuffer,};use futures::{channel::mpsc::UnboundedSender, StreamExt as _};use gpui::{prelude::*, AnyElement, AsyncWindowContext, Model, Task, View, WeakView};use language::ToPoint;use project::{search::SearchQuery, Project, ProjectPath};use schemars::JsonSchema;use serde::Deserialize;use std::path::Path;use ui::prelude::*;use util::ResultExt;use workspace::Workspace; pub struct AnnotationTool {    workspace: WeakView<Workspace>,    project: Model<Project>,} impl AnnotationTool {    pub fn new(workspace: WeakView<Workspace>, project: Model<Project>) -> Self {        Self { workspace, project }    }} #[derive(Default, Debug, Deserialize, JsonSchema, Clone)]pub struct AnnotationInput {    /// Name for this set of annotations    #[serde(default = "default_title")]    title: String,    /// Excerpts from the file to show to the user.    excerpts: Vec<Excerpt>,} fn default_title() -> String {    "Untitled".to_string()} #[derive(Debug, Deserialize, JsonSchema, Clone)]struct Excerpt {    /// Path to the file    path: String,    /// A short, distinctive string that appears in the file, used to define a location in the file.    text_passage: String,    /// Text to display above the code excerpt    annotation: String,} impl LanguageModelTool for AnnotationTool {    type View = AnnotationResultView;     fn name(&self) -> String {        "annotate_code".to_string()    }     fn description(&self) -> String {        "Dynamically annotate symbols in the current codebase. Opens a buffer in a panel in their editor, to the side of the conversation. The annotations are shown in the editor as a block decoration.".to_string()    }     fn view(&self, cx: &mut WindowContext) -> View<Self::View> {        cx.new_view(|cx| {            let (tx, mut rx) = futures::channel::mpsc::unbounded();            cx.spawn(|view, mut cx| async move {                while let Some(excerpt) = rx.next().await {                    AnnotationResultView::add_excerpt(view.clone(), excerpt, &mut cx).await?;                }                anyhow::Ok(())            })            .detach();             AnnotationResultView {                project: self.project.clone(),                workspace: self.workspace.clone(),                tx,                pending_excerpt: None,                added_editor_to_workspace: false,                editor: None,                error: None,                rendered_excerpt_count: 0,            }        })    }} pub struct AnnotationResultView {    workspace: WeakView<Workspace>,    project: Model<Project>,    pending_excerpt: Option<Excerpt>,    added_editor_to_workspace: bool,    editor: Option<View<Editor>>,    tx: UnboundedSender<Excerpt>,    error: Option<anyhow::Error>,    rendered_excerpt_count: usize,} impl AnnotationResultView {    async fn add_excerpt(        this: WeakView<Self>,        excerpt: Excerpt,        cx: &mut AsyncWindowContext,    ) -> Result<()> {        let project = this.update(cx, |this, _cx| this.project.clone())?;         let worktree_id = project.update(cx, |project, cx| {            let worktree = project.worktrees().next()?;            let worktree_id = worktree.read(cx).id();            Some(worktree_id)        })?;         let worktree_id = if let Some(worktree_id) = worktree_id {            worktree_id        } else {            return Err(anyhow::anyhow!("No worktree found"));        };         let buffer_task = project.update(cx, |project, cx| {            project.open_buffer(                ProjectPath {                    worktree_id,                    path: Path::new(&excerpt.path).into(),                },                cx,            )        })?;         let buffer = match buffer_task.await {            Ok(buffer) => buffer,            Err(error) => {                return this.update(cx, |this, cx| {                    this.error = Some(error);                    cx.notify();                })            }        };         let snapshot = buffer.update(cx, |buffer, _cx| buffer.snapshot())?;        let query = SearchQuery::text(&excerpt.text_passage, false, false, false, vec![], vec![])?;        let matches = query.search(&snapshot, None).await;        let Some(first_match) = matches.first() else {            log::warn!(                "text {:?} does not appear in '{}'",                excerpt.text_passage,                excerpt.path            );            return Ok(());        };         this.update(cx, |this, cx| {            let mut start = first_match.start.to_point(&snapshot);            start.column = 0;             if let Some(editor) = &this.editor {                editor.update(cx, |editor, cx| {                    let ranges = editor.buffer().update(cx, |multibuffer, cx| {                        multibuffer.push_excerpts_with_context_lines(                            buffer.clone(),                            vec![start..start],                            5,                            cx,                        )                    });                     let annotation = SharedString::from(excerpt.annotation);                    editor.insert_blocks(                        [BlockProperties {                            position: ranges[0].start,                            height: annotation.split('\n').count() as u8 + 1,                            style: BlockStyle::Fixed,                            render: Box::new(move |cx| Self::render_note_block(&annotation, cx)),                            disposition: BlockDisposition::Above,                        }],                        None,                        cx,                    );                });                 if !this.added_editor_to_workspace {                    this.added_editor_to_workspace = true;                    this.workspace                        .update(cx, |workspace, cx| {                            workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);                        })                        .log_err();                }            }        })?;         Ok(())    }     fn render_note_block(explanation: &SharedString, cx: &mut BlockContext) -> AnyElement {        let anchor_x = cx.anchor_x;        let gutter_width = cx.gutter_dimensions.width;         h_flex()            .w_full()            .py_2()            .border_y_1()            .border_color(cx.theme().colors().border)            .child(                h_flex()                    .justify_center()                    .w(gutter_width)                    .child(Icon::new(IconName::Ai).color(Color::Hint)),            )            .child(                h_flex()                    .w_full()                    .ml(anchor_x - gutter_width)                    .child(explanation.clone()),            )            .into_any_element()    }} impl Render for AnnotationResultView {    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {        if let Some(error) = &self.error {            ui::Label::new(error.to_string()).into_any_element()        } else {            ui::Label::new(SharedString::from(format!(                "Opened a buffer with {} excerpts",                self.rendered_excerpt_count            )))            .into_any_element()        }    }} impl ToolView for AnnotationResultView {    type Input = AnnotationInput;    type SerializedState = Option<String>;     fn generate(&self, _: &mut ProjectContext, _: &mut ViewContext<Self>) -> String {        if let Some(error) = &self.error {            format!("Failed to create buffer: {error:?}")        } else {            format!(                "opened {} excerpts in a buffer",                self.rendered_excerpt_count            )        }    }     fn set_input(&mut self, mut input: Self::Input, cx: &mut ViewContext<Self>) {        let editor = if let Some(editor) = &self.editor {            editor.clone()        } else {            let multibuffer = cx.new_model(|_cx| {                MultiBuffer::new(0, language::Capability::ReadWrite).with_title(String::new())            });            let editor = cx.new_view(|cx| {                Editor::for_multibuffer(multibuffer.clone(), Some(self.project.clone()), cx)            });             self.editor = Some(editor.clone());            editor        };         editor.update(cx, |editor, cx| {            editor.buffer().update(cx, |multibuffer, cx| {                if multibuffer.title(cx) != input.title {                    multibuffer.set_title(input.title.clone(), cx);                }            });             self.pending_excerpt = input.excerpts.pop();            for excerpt in input.excerpts.iter().skip(self.rendered_excerpt_count) {                self.tx.unbounded_send(excerpt.clone()).ok();            }            self.rendered_excerpt_count = input.excerpts.len();        });         cx.notify();    }     fn execute(&mut self, _cx: &mut ViewContext<Self>) -> Task<Result<()>> {        if let Some(excerpt) = self.pending_excerpt.take() {            self.rendered_excerpt_count += 1;            self.tx.unbounded_send(excerpt.clone()).ok();        }         self.tx.close_channel();        Task::ready(Ok(()))    }     fn serialize(&self, _cx: &mut ViewContext<Self>) -> Self::SerializedState {        self.error.as_ref().map(|error| error.to_string())    }     fn deserialize(        &mut self,        output: Self::SerializedState,        _cx: &mut ViewContext<Self>,    ) -> Result<()> {        if let Some(error_message) = output {            self.error = Some(anyhow::anyhow!("{}", error_message));        }        Ok(())    }}