1use anyhow::Result;
2use assistant_tooling::{LanguageModelTool, ProjectContext, ToolView};
3use editor::{
4 display_map::{BlockContext, BlockDisposition, BlockProperties, BlockStyle},
5 Editor, MultiBuffer,
6};
7use futures::{channel::mpsc::UnboundedSender, StreamExt as _};
8use gpui::{prelude::*, AnyElement, AsyncWindowContext, Model, Task, View, WeakView};
9use language::ToPoint;
10use project::{search::SearchQuery, Project, ProjectPath};
11use schemars::JsonSchema;
12use serde::Deserialize;
13use std::path::Path;
14use ui::prelude::*;
15use util::ResultExt;
16use workspace::Workspace;
17
18pub struct AnnotationTool {
19 workspace: WeakView<Workspace>,
20 project: Model<Project>,
21}
22
23impl AnnotationTool {
24 pub fn new(workspace: WeakView<Workspace>, project: Model<Project>) -> Self {
25 Self { workspace, project }
26 }
27}
28
29#[derive(Default, Debug, Deserialize, JsonSchema, Clone)]
30pub struct AnnotationInput {
31 /// Name for this set of annotations
32 #[serde(default = "default_title")]
33 title: String,
34 /// Excerpts from the file to show to the user.
35 excerpts: Vec<Excerpt>,
36}
37
38fn default_title() -> String {
39 "Untitled".to_string()
40}
41
42#[derive(Debug, Deserialize, JsonSchema, Clone)]
43struct Excerpt {
44 /// Path to the file
45 path: String,
46 /// A short, distinctive string that appears in the file, used to define a location in the file.
47 text_passage: String,
48 /// Text to display above the code excerpt
49 annotation: String,
50}
51
52impl LanguageModelTool for AnnotationTool {
53 type View = AnnotationResultView;
54
55 fn name(&self) -> String {
56 "annotate_code".to_string()
57 }
58
59 fn description(&self) -> String {
60 "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()
61 }
62
63 fn view(&self, cx: &mut WindowContext) -> View<Self::View> {
64 cx.new_view(|cx| {
65 let (tx, mut rx) = futures::channel::mpsc::unbounded();
66 cx.spawn(|view, mut cx| async move {
67 while let Some(excerpt) = rx.next().await {
68 AnnotationResultView::add_excerpt(view.clone(), excerpt, &mut cx).await?;
69 }
70 anyhow::Ok(())
71 })
72 .detach();
73
74 AnnotationResultView {
75 project: self.project.clone(),
76 workspace: self.workspace.clone(),
77 tx,
78 pending_excerpt: None,
79 added_editor_to_workspace: false,
80 editor: None,
81 error: None,
82 rendered_excerpt_count: 0,
83 }
84 })
85 }
86}
87
88pub struct AnnotationResultView {
89 workspace: WeakView<Workspace>,
90 project: Model<Project>,
91 pending_excerpt: Option<Excerpt>,
92 added_editor_to_workspace: bool,
93 editor: Option<View<Editor>>,
94 tx: UnboundedSender<Excerpt>,
95 error: Option<anyhow::Error>,
96 rendered_excerpt_count: usize,
97}
98
99impl AnnotationResultView {
100 async fn add_excerpt(
101 this: WeakView<Self>,
102 excerpt: Excerpt,
103 cx: &mut AsyncWindowContext,
104 ) -> Result<()> {
105 let project = this.update(cx, |this, _cx| this.project.clone())?;
106
107 let worktree_id = project.update(cx, |project, cx| {
108 let worktree = project.worktrees().next()?;
109 let worktree_id = worktree.read(cx).id();
110 Some(worktree_id)
111 })?;
112
113 let worktree_id = if let Some(worktree_id) = worktree_id {
114 worktree_id
115 } else {
116 return Err(anyhow::anyhow!("No worktree found"));
117 };
118
119 let buffer_task = project.update(cx, |project, cx| {
120 project.open_buffer(
121 ProjectPath {
122 worktree_id,
123 path: Path::new(&excerpt.path).into(),
124 },
125 cx,
126 )
127 })?;
128
129 let buffer = match buffer_task.await {
130 Ok(buffer) => buffer,
131 Err(error) => {
132 return this.update(cx, |this, cx| {
133 this.error = Some(error);
134 cx.notify();
135 })
136 }
137 };
138
139 let snapshot = buffer.update(cx, |buffer, _cx| buffer.snapshot())?;
140 let query = SearchQuery::text(&excerpt.text_passage, false, false, false, vec![], vec![])?;
141 let matches = query.search(&snapshot, None).await;
142 let Some(first_match) = matches.first() else {
143 log::warn!(
144 "text {:?} does not appear in '{}'",
145 excerpt.text_passage,
146 excerpt.path
147 );
148 return Ok(());
149 };
150
151 this.update(cx, |this, cx| {
152 let mut start = first_match.start.to_point(&snapshot);
153 start.column = 0;
154
155 if let Some(editor) = &this.editor {
156 editor.update(cx, |editor, cx| {
157 let ranges = editor.buffer().update(cx, |multibuffer, cx| {
158 multibuffer.push_excerpts_with_context_lines(
159 buffer.clone(),
160 vec![start..start],
161 5,
162 cx,
163 )
164 });
165
166 let annotation = SharedString::from(excerpt.annotation);
167 editor.insert_blocks(
168 [BlockProperties {
169 position: ranges[0].start,
170 height: annotation.split('\n').count() as u8 + 1,
171 style: BlockStyle::Fixed,
172 render: Box::new(move |cx| Self::render_note_block(&annotation, cx)),
173 disposition: BlockDisposition::Above,
174 }],
175 None,
176 cx,
177 );
178 });
179
180 if !this.added_editor_to_workspace {
181 this.added_editor_to_workspace = true;
182 this.workspace
183 .update(cx, |workspace, cx| {
184 workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);
185 })
186 .log_err();
187 }
188 }
189 })?;
190
191 Ok(())
192 }
193
194 fn render_note_block(explanation: &SharedString, cx: &mut BlockContext) -> AnyElement {
195 let anchor_x = cx.anchor_x;
196 let gutter_width = cx.gutter_dimensions.width;
197
198 h_flex()
199 .w_full()
200 .py_2()
201 .border_y_1()
202 .border_color(cx.theme().colors().border)
203 .child(
204 h_flex()
205 .justify_center()
206 .w(gutter_width)
207 .child(Icon::new(IconName::Ai).color(Color::Hint)),
208 )
209 .child(
210 h_flex()
211 .w_full()
212 .ml(anchor_x - gutter_width)
213 .child(explanation.clone()),
214 )
215 .into_any_element()
216 }
217}
218
219impl Render for AnnotationResultView {
220 fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
221 if let Some(error) = &self.error {
222 ui::Label::new(error.to_string()).into_any_element()
223 } else {
224 ui::Label::new(SharedString::from(format!(
225 "Opened a buffer with {} excerpts",
226 self.rendered_excerpt_count
227 )))
228 .into_any_element()
229 }
230 }
231}
232
233impl ToolView for AnnotationResultView {
234 type Input = AnnotationInput;
235 type SerializedState = Option<String>;
236
237 fn generate(&self, _: &mut ProjectContext, _: &mut ViewContext<Self>) -> String {
238 if let Some(error) = &self.error {
239 format!("Failed to create buffer: {error:?}")
240 } else {
241 format!(
242 "opened {} excerpts in a buffer",
243 self.rendered_excerpt_count
244 )
245 }
246 }
247
248 fn set_input(&mut self, mut input: Self::Input, cx: &mut ViewContext<Self>) {
249 let editor = if let Some(editor) = &self.editor {
250 editor.clone()
251 } else {
252 let multibuffer = cx.new_model(|_cx| {
253 MultiBuffer::new(0, language::Capability::ReadWrite).with_title(String::new())
254 });
255 let editor = cx.new_view(|cx| {
256 Editor::for_multibuffer(multibuffer.clone(), Some(self.project.clone()), cx)
257 });
258
259 self.editor = Some(editor.clone());
260 editor
261 };
262
263 editor.update(cx, |editor, cx| {
264 editor.buffer().update(cx, |multibuffer, cx| {
265 if multibuffer.title(cx) != input.title {
266 multibuffer.set_title(input.title.clone(), cx);
267 }
268 });
269
270 self.pending_excerpt = input.excerpts.pop();
271 for excerpt in input.excerpts.iter().skip(self.rendered_excerpt_count) {
272 self.tx.unbounded_send(excerpt.clone()).ok();
273 }
274 self.rendered_excerpt_count = input.excerpts.len();
275 });
276
277 cx.notify();
278 }
279
280 fn execute(&mut self, _cx: &mut ViewContext<Self>) -> Task<Result<()>> {
281 if let Some(excerpt) = self.pending_excerpt.take() {
282 self.rendered_excerpt_count += 1;
283 self.tx.unbounded_send(excerpt.clone()).ok();
284 }
285
286 self.tx.close_channel();
287 Task::ready(Ok(()))
288 }
289
290 fn serialize(&self, _cx: &mut ViewContext<Self>) -> Self::SerializedState {
291 self.error.as_ref().map(|error| error.to_string())
292 }
293
294 fn deserialize(
295 &mut self,
296 output: Self::SerializedState,
297 _cx: &mut ViewContext<Self>,
298 ) -> Result<()> {
299 if let Some(error_message) = output {
300 self.error = Some(anyhow::anyhow!("{}", error_message));
301 }
302 Ok(())
303 }
304}