1use anyhow::Result;
2use assistant_tooling::LanguageModelTool;
3use gpui::{prelude::*, AnyElement, AppContext, Model, Task};
4use project::Fs;
5use schemars::JsonSchema;
6use semantic_index::ProjectIndex;
7use serde::{Deserialize, Serialize};
8use std::sync::Arc;
9use ui::{
10 div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, SharedString,
11 WindowContext,
12};
13use util::ResultExt as _;
14
15const DEFAULT_SEARCH_LIMIT: usize = 20;
16
17#[derive(Serialize, Clone)]
18pub struct CodebaseExcerpt {
19 path: SharedString,
20 text: SharedString,
21 score: f32,
22}
23
24// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model.
25// Any changes or deletions to the `CodebaseQuery` comments will change model behavior.
26
27#[derive(Deserialize, JsonSchema)]
28pub struct CodebaseQuery {
29 /// Semantic search query
30 query: String,
31 /// Maximum number of results to return, defaults to 20
32 limit: Option<usize>,
33}
34
35pub struct ProjectIndexTool {
36 project_index: Model<ProjectIndex>,
37 fs: Arc<dyn Fs>,
38}
39
40impl ProjectIndexTool {
41 pub fn new(project_index: Model<ProjectIndex>, fs: Arc<dyn Fs>) -> Self {
42 // TODO: setup a better description based on the user's current codebase.
43 Self { project_index, fs }
44 }
45}
46
47impl LanguageModelTool for ProjectIndexTool {
48 type Input = CodebaseQuery;
49 type Output = Vec<CodebaseExcerpt>;
50
51 fn name(&self) -> String {
52 "query_codebase".to_string()
53 }
54
55 fn description(&self) -> String {
56 "Semantic search against the user's current codebase, returning excerpts related to the query by computing a dot product against embeddings of chunks and an embedding of the query".to_string()
57 }
58
59 fn execute(&self, query: &Self::Input, cx: &AppContext) -> Task<Result<Self::Output>> {
60 let project_index = self.project_index.read(cx);
61
62 let results = project_index.search(
63 query.query.as_str(),
64 query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
65 cx,
66 );
67
68 let fs = self.fs.clone();
69
70 cx.spawn(|cx| async move {
71 let results = results.await;
72
73 let excerpts = results.into_iter().map(|result| {
74 let abs_path = result
75 .worktree
76 .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
77 let fs = fs.clone();
78
79 async move {
80 let path = result.path.clone();
81 let text = fs.load(&abs_path?).await?;
82
83 let mut start = result.range.start;
84 let mut end = result.range.end.min(text.len());
85 while !text.is_char_boundary(start) {
86 start += 1;
87 }
88 while !text.is_char_boundary(end) {
89 end -= 1;
90 }
91
92 anyhow::Ok(CodebaseExcerpt {
93 path: path.to_string_lossy().to_string().into(),
94 text: SharedString::from(text[start..end].to_string()),
95 score: result.score,
96 })
97 }
98 });
99
100 let excerpts = futures::future::join_all(excerpts)
101 .await
102 .into_iter()
103 .filter_map(|result| result.log_err())
104 .collect();
105 anyhow::Ok(excerpts)
106 })
107 }
108
109 fn render(
110 _tool_call_id: &str,
111 input: &Self::Input,
112 excerpts: &Self::Output,
113 cx: &mut WindowContext,
114 ) -> AnyElement {
115 let query = input.query.clone();
116
117 div()
118 .v_flex()
119 .gap_2()
120 .child(
121 div()
122 .p_2()
123 .rounded_md()
124 .bg(cx.theme().colors().editor_background)
125 .child(
126 h_flex()
127 .child(Label::new("Query: ").color(Color::Modified))
128 .child(Label::new(query).color(Color::Muted)),
129 ),
130 )
131 .children(excerpts.iter().map(|excerpt| {
132 // This render doesn't have state/model, so we can't use the listener
133 // let expanded = excerpt.expanded;
134 // let element_id = excerpt.element_id.clone();
135 let element_id = ElementId::Name(nanoid::nanoid!().into());
136 let expanded = false;
137
138 CollapsibleContainer::new(element_id.clone(), expanded)
139 .start_slot(
140 h_flex()
141 .gap_1()
142 .child(Icon::new(IconName::File).color(Color::Muted))
143 .child(Label::new(excerpt.path.clone()).color(Color::Muted)),
144 )
145 // .on_click(cx.listener(move |this, _, cx| {
146 // this.toggle_expanded(element_id.clone(), cx);
147 // }))
148 .child(
149 div()
150 .p_2()
151 .rounded_md()
152 .bg(cx.theme().colors().editor_background)
153 .child(
154 excerpt.text.clone(), // todo!(): Show as an editor block
155 ),
156 )
157 }))
158 .into_any_element()
159 }
160
161 fn format(_input: &Self::Input, excerpts: &Self::Output) -> String {
162 let mut body = "Semantic search results:\n".to_string();
163
164 for excerpt in excerpts {
165 body.push_str("Excerpt from ");
166 body.push_str(excerpt.path.as_ref());
167 body.push_str(", score ");
168 body.push_str(&excerpt.score.to_string());
169 body.push_str(":\n");
170 body.push_str("~~~\n");
171 body.push_str(excerpt.text.as_ref());
172 body.push_str("~~~\n");
173 }
174 body
175 }
176}