1use anyhow::Result;
2use assistant_tooling::{LanguageModelTool, ProjectContext, ToolOutput};
3use editor::{
4 display_map::{BlockContext, BlockDisposition, BlockProperties, BlockStyle},
5 Editor, MultiBuffer,
6};
7use gpui::{prelude::*, AnyElement, Model, Task, View, WeakView};
8use language::ToPoint;
9use project::{search::SearchQuery, Project, ProjectPath};
10use schemars::JsonSchema;
11use serde::Deserialize;
12use std::path::Path;
13use ui::prelude::*;
14use util::ResultExt;
15use workspace::Workspace;
16
17pub struct AnnotationTool {
18 workspace: WeakView<Workspace>,
19 project: Model<Project>,
20}
21
22impl AnnotationTool {
23 pub fn new(workspace: WeakView<Workspace>, project: Model<Project>) -> Self {
24 Self { workspace, project }
25 }
26}
27
28#[derive(Debug, Deserialize, JsonSchema, Clone)]
29pub struct AnnotationInput {
30 /// Name for this set of annotations
31 title: String,
32 /// Excerpts from the file to show to the user.
33 excerpts: Vec<Excerpt>,
34}
35
36#[derive(Debug, Deserialize, JsonSchema, Clone)]
37struct Excerpt {
38 /// Path to the file
39 path: String,
40 /// A short, distinctive string that appears in the file, used to define a location in the file.
41 text_passage: String,
42 /// Text to display above the code excerpt
43 annotation: String,
44}
45
46impl LanguageModelTool for AnnotationTool {
47 type Input = AnnotationInput;
48 type Output = String;
49 type View = AnnotationResultView;
50
51 fn name(&self) -> String {
52 "annotate_code".to_string()
53 }
54
55 fn description(&self) -> String {
56 "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()
57 }
58
59 fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
60 let workspace = self.workspace.clone();
61 let project = self.project.clone();
62 let excerpts = input.excerpts.clone();
63 let title = input.title.clone();
64
65 let worktree_id = project.update(cx, |project, cx| {
66 let worktree = project.worktrees().next()?;
67 let worktree_id = worktree.read(cx).id();
68 Some(worktree_id)
69 });
70
71 let worktree_id = if let Some(worktree_id) = worktree_id {
72 worktree_id
73 } else {
74 return Task::ready(Err(anyhow::anyhow!("No worktree found")));
75 };
76
77 let buffer_tasks = project.update(cx, |project, cx| {
78 excerpts
79 .iter()
80 .map(|excerpt| {
81 project.open_buffer(
82 ProjectPath {
83 worktree_id,
84 path: Path::new(&excerpt.path).into(),
85 },
86 cx,
87 )
88 })
89 .collect::<Vec<_>>()
90 });
91
92 cx.spawn(move |mut cx| async move {
93 let buffers = futures::future::try_join_all(buffer_tasks).await?;
94
95 let multibuffer = cx.new_model(|_cx| {
96 MultiBuffer::new(0, language::Capability::ReadWrite).with_title(title)
97 })?;
98 let editor =
99 cx.new_view(|cx| Editor::for_multibuffer(multibuffer, Some(project), cx))?;
100
101 for (excerpt, buffer) in excerpts.iter().zip(buffers.iter()) {
102 let snapshot = buffer.update(&mut cx, |buffer, _cx| buffer.snapshot())?;
103
104 let query =
105 SearchQuery::text(&excerpt.text_passage, false, false, false, vec![], vec![])?;
106
107 let matches = query.search(&snapshot, None).await;
108 let Some(first_match) = matches.first() else {
109 log::warn!(
110 "text {:?} does not appear in '{}'",
111 excerpt.text_passage,
112 excerpt.path
113 );
114 continue;
115 };
116 let mut start = first_match.start.to_point(&snapshot);
117 start.column = 0;
118
119 editor.update(&mut cx, |editor, cx| {
120 let ranges = editor.buffer().update(cx, |multibuffer, cx| {
121 multibuffer.push_excerpts_with_context_lines(
122 buffer.clone(),
123 vec![start..start],
124 5,
125 cx,
126 )
127 });
128 let annotation = SharedString::from(excerpt.annotation.clone());
129 editor.insert_blocks(
130 [BlockProperties {
131 position: ranges[0].start,
132 height: annotation.split('\n').count() as u8 + 1,
133 style: BlockStyle::Fixed,
134 render: Box::new(move |cx| Self::render_note_block(&annotation, cx)),
135 disposition: BlockDisposition::Above,
136 }],
137 None,
138 cx,
139 );
140 })?;
141 }
142
143 workspace
144 .update(&mut cx, |workspace, cx| {
145 workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);
146 })
147 .log_err();
148
149 anyhow::Ok("showed comments to users in a new view".into())
150 })
151 }
152
153 fn view(
154 &self,
155 _: Self::Input,
156 output: Result<Self::Output>,
157 cx: &mut WindowContext,
158 ) -> View<Self::View> {
159 cx.new_view(|_cx| AnnotationResultView { output })
160 }
161}
162
163impl AnnotationTool {
164 fn render_note_block(explanation: &SharedString, cx: &mut BlockContext) -> AnyElement {
165 let anchor_x = cx.anchor_x;
166 let gutter_width = cx.gutter_dimensions.width;
167
168 h_flex()
169 .w_full()
170 .py_2()
171 .border_y_1()
172 .border_color(cx.theme().colors().border)
173 .child(
174 h_flex()
175 .justify_center()
176 .w(gutter_width)
177 .child(Icon::new(IconName::Ai).color(Color::Hint)),
178 )
179 .child(
180 h_flex()
181 .w_full()
182 .ml(anchor_x - gutter_width)
183 .child(explanation.clone()),
184 )
185 .into_any_element()
186 }
187}
188
189pub struct AnnotationResultView {
190 output: Result<String>,
191}
192
193impl Render for AnnotationResultView {
194 fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
195 match &self.output {
196 Ok(output) => div().child(output.clone().into_any_element()),
197 Err(error) => div().child(format!("failed to open path: {:?}", error)),
198 }
199 }
200}
201
202impl ToolOutput for AnnotationResultView {
203 fn generate(&self, _: &mut ProjectContext, _: &mut WindowContext) -> String {
204 match &self.output {
205 Ok(output) => output.clone(),
206 Err(err) => format!("Failed to create buffer: {err:?}"),
207 }
208 }
209}