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::{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 annotations: Vec<Annotation>,
33}
34
35#[derive(Debug, Deserialize, JsonSchema, Clone)]
36struct Annotation {
37 /// Path to the file
38 path: String,
39 /// Name of a symbol in the code
40 symbol_name: String,
41 /// Text to display near the symbol definition
42 text: String,
43}
44
45impl LanguageModelTool for AnnotationTool {
46 type Input = AnnotationInput;
47 type Output = String;
48 type View = AnnotationResultView;
49
50 fn name(&self) -> String {
51 "annotate_code".to_string()
52 }
53
54 fn description(&self) -> String {
55 "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()
56 }
57
58 fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
59 let workspace = self.workspace.clone();
60 let project = self.project.clone();
61 let excerpts = input.annotations.clone();
62 let title = input.title.clone();
63
64 let worktree_id = project.update(cx, |project, cx| {
65 let worktree = project.worktrees().next()?;
66 let worktree_id = worktree.read(cx).id();
67 Some(worktree_id)
68 });
69
70 let worktree_id = if let Some(worktree_id) = worktree_id {
71 worktree_id
72 } else {
73 return Task::ready(Err(anyhow::anyhow!("No worktree found")));
74 };
75
76 let buffer_tasks = project.update(cx, |project, cx| {
77 let excerpts = excerpts.clone();
78 excerpts
79 .iter()
80 .map(|excerpt| {
81 let project_path = ProjectPath {
82 worktree_id,
83 path: Path::new(&excerpt.path).into(),
84 };
85 project.open_buffer(project_path.clone(), cx)
86 })
87 .collect::<Vec<_>>()
88 });
89
90 cx.spawn(move |mut cx| async move {
91 let buffers = futures::future::try_join_all(buffer_tasks).await?;
92
93 let multibuffer = cx.new_model(|_cx| {
94 MultiBuffer::new(0, language::Capability::ReadWrite).with_title(title)
95 })?;
96 let editor =
97 cx.new_view(|cx| Editor::for_multibuffer(multibuffer, Some(project), cx))?;
98
99 for (excerpt, buffer) in excerpts.iter().zip(buffers.iter()) {
100 let snapshot = buffer.update(&mut cx, |buffer, _cx| buffer.snapshot())?;
101
102 if let Some(outline) = snapshot.outline(None) {
103 let matches = outline
104 .search(&excerpt.symbol_name, cx.background_executor().clone())
105 .await;
106 if let Some(mat) = matches.first() {
107 let item = &outline.items[mat.candidate_id];
108 let start = item.range.start.to_point(&snapshot);
109 editor.update(&mut cx, |editor, cx| {
110 let ranges = editor.buffer().update(cx, |multibuffer, cx| {
111 multibuffer.push_excerpts_with_context_lines(
112 buffer.clone(),
113 vec![start..start],
114 5,
115 cx,
116 )
117 });
118 let explanation = SharedString::from(excerpt.text.clone());
119 editor.insert_blocks(
120 [BlockProperties {
121 position: ranges[0].start,
122 height: 2,
123 style: BlockStyle::Fixed,
124 render: Box::new(move |cx| {
125 Self::render_note_block(&explanation, cx)
126 }),
127 disposition: BlockDisposition::Above,
128 }],
129 None,
130 cx,
131 );
132 })?;
133 }
134 }
135 }
136
137 workspace
138 .update(&mut cx, |workspace, cx| {
139 workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);
140 })
141 .log_err();
142
143 anyhow::Ok("showed comments to users in a new view".into())
144 })
145 }
146
147 fn output_view(
148 _: Self::Input,
149 output: Result<Self::Output>,
150 cx: &mut WindowContext,
151 ) -> View<Self::View> {
152 cx.new_view(|_cx| AnnotationResultView { output })
153 }
154}
155
156impl AnnotationTool {
157 fn render_note_block(explanation: &SharedString, cx: &mut BlockContext) -> AnyElement {
158 let anchor_x = cx.anchor_x;
159 let gutter_width = cx.gutter_dimensions.width;
160
161 h_flex()
162 .w_full()
163 .py_2()
164 .border_y_1()
165 .border_color(cx.theme().colors().border)
166 .child(
167 h_flex()
168 .justify_center()
169 .w(gutter_width)
170 .child(Icon::new(IconName::Ai).color(Color::Hint)),
171 )
172 .child(
173 h_flex()
174 .w_full()
175 .ml(anchor_x - gutter_width)
176 .child(explanation.clone()),
177 )
178 .into_any_element()
179 }
180}
181
182pub struct AnnotationResultView {
183 output: Result<String>,
184}
185
186impl Render for AnnotationResultView {
187 fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
188 match &self.output {
189 Ok(output) => div().child(output.clone().into_any_element()),
190 Err(error) => div().child(format!("failed to open path: {:?}", error)),
191 }
192 }
193}
194
195impl ToolOutput for AnnotationResultView {
196 fn generate(&self, _: &mut ProjectContext, _: &mut WindowContext) -> String {
197 match &self.output {
198 Ok(output) => output.clone(),
199 Err(err) => format!("Failed to create buffer: {err:?}"),
200 }
201 }
202}