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(()) }}
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(())
}
}