1use std::ops::Range;
2
3use file_icons::FileIcons;
4use gpui::{App, Entity, SharedString};
5use language::Buffer;
6use language_model::{LanguageModelRequestMessage, MessageContent};
7use project::ProjectPath;
8use serde::{Deserialize, Serialize};
9use text::{Anchor, BufferId};
10use ui::IconName;
11use util::post_inc;
12
13use crate::{context_store::buffer_path_log_err, thread::Thread};
14
15#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
16pub struct ContextId(pub(crate) usize);
17
18impl ContextId {
19 pub fn post_inc(&mut self) -> Self {
20 Self(post_inc(&mut self.0))
21 }
22}
23
24/// Some context attached to a message in a thread.
25#[derive(Debug, Clone)]
26pub struct ContextSnapshot {
27 pub id: ContextId,
28 pub name: SharedString,
29 pub parent: Option<SharedString>,
30 pub tooltip: Option<SharedString>,
31 pub icon_path: Option<SharedString>,
32 pub kind: ContextKind,
33 /// Joining these strings separated by \n yields text for model. Not refreshed by `snapshot`.
34 pub text: Box<[SharedString]>,
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum ContextKind {
39 File,
40 Directory,
41 Symbol,
42 FetchedUrl,
43 Thread,
44}
45
46impl ContextKind {
47 pub fn icon(&self) -> IconName {
48 match self {
49 ContextKind::File => IconName::File,
50 ContextKind::Directory => IconName::Folder,
51 ContextKind::Symbol => IconName::Code,
52 ContextKind::FetchedUrl => IconName::Globe,
53 ContextKind::Thread => IconName::MessageCircle,
54 }
55 }
56}
57
58#[derive(Debug)]
59pub enum AssistantContext {
60 File(FileContext),
61 Directory(DirectoryContext),
62 Symbol(SymbolContext),
63 FetchedUrl(FetchedUrlContext),
64 Thread(ThreadContext),
65}
66
67impl AssistantContext {
68 pub fn id(&self) -> ContextId {
69 match self {
70 Self::File(file) => file.id,
71 Self::Directory(directory) => directory.snapshot.id,
72 Self::Symbol(symbol) => symbol.id,
73 Self::FetchedUrl(url) => url.id,
74 Self::Thread(thread) => thread.id,
75 }
76 }
77}
78
79#[derive(Debug)]
80pub struct FileContext {
81 pub id: ContextId,
82 pub context_buffer: ContextBuffer,
83}
84
85#[derive(Debug)]
86pub struct DirectoryContext {
87 pub path: ProjectPath,
88 pub context_buffers: Vec<ContextBuffer>,
89 pub snapshot: ContextSnapshot,
90}
91
92#[derive(Debug)]
93pub struct SymbolContext {
94 pub id: ContextId,
95 pub context_symbol: ContextSymbol,
96}
97
98#[derive(Debug)]
99pub struct FetchedUrlContext {
100 pub id: ContextId,
101 pub url: SharedString,
102 pub text: SharedString,
103}
104
105// TODO: Model<Thread> holds onto the thread even if the thread is deleted. Can either handle this
106// explicitly or have a WeakModel<Thread> and remove during snapshot.
107
108#[derive(Debug)]
109pub struct ThreadContext {
110 pub id: ContextId,
111 pub thread: Entity<Thread>,
112 pub text: SharedString,
113}
114
115// TODO: Model<Buffer> holds onto the buffer even if the file is deleted and closed. Should remove
116// the context from the message editor in this case.
117
118#[derive(Debug, Clone)]
119pub struct ContextBuffer {
120 pub id: BufferId,
121 pub buffer: Entity<Buffer>,
122 pub version: clock::Global,
123 pub text: SharedString,
124}
125
126#[derive(Debug, Clone)]
127pub struct ContextSymbol {
128 pub id: ContextSymbolId,
129 pub buffer: Entity<Buffer>,
130 pub buffer_version: clock::Global,
131 /// The range that the symbol encloses, e.g. for function symbol, this will
132 /// include not only the signature, but also the body
133 pub enclosing_range: Range<Anchor>,
134 pub text: SharedString,
135}
136
137#[derive(Debug, Clone, PartialEq, Eq, Hash)]
138pub struct ContextSymbolId {
139 pub path: ProjectPath,
140 pub name: SharedString,
141 pub range: Range<Anchor>,
142}
143
144impl AssistantContext {
145 pub fn snapshot(&self, cx: &App) -> Option<ContextSnapshot> {
146 match &self {
147 Self::File(file_context) => file_context.snapshot(cx),
148 Self::Directory(directory_context) => Some(directory_context.snapshot()),
149 Self::Symbol(symbol_context) => symbol_context.snapshot(cx),
150 Self::FetchedUrl(fetched_url_context) => Some(fetched_url_context.snapshot()),
151 Self::Thread(thread_context) => Some(thread_context.snapshot(cx)),
152 }
153 }
154}
155
156impl FileContext {
157 pub fn snapshot(&self, cx: &App) -> Option<ContextSnapshot> {
158 let buffer = self.context_buffer.buffer.read(cx);
159 let path = buffer_path_log_err(buffer)?;
160 let full_path: SharedString = path.to_string_lossy().into_owned().into();
161 let name = match path.file_name() {
162 Some(name) => name.to_string_lossy().into_owned().into(),
163 None => full_path.clone(),
164 };
165 let parent = path
166 .parent()
167 .and_then(|p| p.file_name())
168 .map(|p| p.to_string_lossy().into_owned().into());
169
170 let icon_path = FileIcons::get_icon(&path, cx);
171
172 Some(ContextSnapshot {
173 id: self.id,
174 name,
175 parent,
176 tooltip: Some(full_path),
177 icon_path,
178 kind: ContextKind::File,
179 text: Box::new([self.context_buffer.text.clone()]),
180 })
181 }
182}
183
184impl DirectoryContext {
185 pub fn new(
186 id: ContextId,
187 project_path: ProjectPath,
188 context_buffers: Vec<ContextBuffer>,
189 ) -> DirectoryContext {
190 let full_path: SharedString = project_path.path.to_string_lossy().into_owned().into();
191
192 let name = match project_path.path.file_name() {
193 Some(name) => name.to_string_lossy().into_owned().into(),
194 None => full_path.clone(),
195 };
196
197 let parent = project_path
198 .path
199 .parent()
200 .and_then(|p| p.file_name())
201 .map(|p| p.to_string_lossy().into_owned().into());
202
203 // TODO: include directory path in text?
204 let text = context_buffers
205 .iter()
206 .map(|b| b.text.clone())
207 .collect::<Vec<_>>()
208 .into();
209
210 DirectoryContext {
211 path: project_path,
212 context_buffers,
213 snapshot: ContextSnapshot {
214 id,
215 name,
216 parent,
217 tooltip: Some(full_path),
218 icon_path: None,
219 kind: ContextKind::Directory,
220 text,
221 },
222 }
223 }
224
225 pub fn snapshot(&self) -> ContextSnapshot {
226 self.snapshot.clone()
227 }
228}
229
230impl SymbolContext {
231 pub fn snapshot(&self, cx: &App) -> Option<ContextSnapshot> {
232 let buffer = self.context_symbol.buffer.read(cx);
233 let name = self.context_symbol.id.name.clone();
234 let path = buffer_path_log_err(buffer)?
235 .to_string_lossy()
236 .into_owned()
237 .into();
238
239 Some(ContextSnapshot {
240 id: self.id,
241 name,
242 parent: Some(path),
243 tooltip: None,
244 icon_path: None,
245 kind: ContextKind::Symbol,
246 text: Box::new([self.context_symbol.text.clone()]),
247 })
248 }
249}
250
251impl FetchedUrlContext {
252 pub fn snapshot(&self) -> ContextSnapshot {
253 ContextSnapshot {
254 id: self.id,
255 name: self.url.clone(),
256 parent: None,
257 tooltip: None,
258 icon_path: None,
259 kind: ContextKind::FetchedUrl,
260 text: Box::new([self.text.clone()]),
261 }
262 }
263}
264
265impl ThreadContext {
266 pub fn snapshot(&self, cx: &App) -> ContextSnapshot {
267 let thread = self.thread.read(cx);
268 ContextSnapshot {
269 id: self.id,
270 name: thread.summary().unwrap_or("New thread".into()),
271 parent: None,
272 tooltip: None,
273 icon_path: None,
274 kind: ContextKind::Thread,
275 text: Box::new([self.text.clone()]),
276 }
277 }
278}
279
280pub fn attach_context_to_message(
281 message: &mut LanguageModelRequestMessage,
282 contexts: impl Iterator<Item = ContextSnapshot>,
283) {
284 let mut file_context = Vec::new();
285 let mut directory_context = Vec::new();
286 let mut symbol_context = Vec::new();
287 let mut fetch_context = Vec::new();
288 let mut thread_context = Vec::new();
289
290 let mut capacity = 0;
291 for context in contexts {
292 capacity += context.text.len();
293 match context.kind {
294 ContextKind::File => file_context.push(context),
295 ContextKind::Directory => directory_context.push(context),
296 ContextKind::Symbol => symbol_context.push(context),
297 ContextKind::FetchedUrl => fetch_context.push(context),
298 ContextKind::Thread => thread_context.push(context),
299 }
300 }
301 if !file_context.is_empty() {
302 capacity += 1;
303 }
304 if !directory_context.is_empty() {
305 capacity += 1;
306 }
307 if !symbol_context.is_empty() {
308 capacity += 1;
309 }
310 if !fetch_context.is_empty() {
311 capacity += 1 + fetch_context.len();
312 }
313 if !thread_context.is_empty() {
314 capacity += 1 + thread_context.len();
315 }
316 if capacity == 0 {
317 return;
318 }
319
320 let mut context_chunks = Vec::with_capacity(capacity);
321
322 if !file_context.is_empty() {
323 context_chunks.push("The following files are available:\n");
324 for context in &file_context {
325 for chunk in &context.text {
326 context_chunks.push(&chunk);
327 }
328 }
329 }
330
331 if !directory_context.is_empty() {
332 context_chunks.push("The following directories are available:\n");
333 for context in &directory_context {
334 for chunk in &context.text {
335 context_chunks.push(&chunk);
336 }
337 }
338 }
339
340 if !symbol_context.is_empty() {
341 context_chunks.push("The following symbols are available:\n");
342 for context in &symbol_context {
343 for chunk in &context.text {
344 context_chunks.push(&chunk);
345 }
346 }
347 }
348
349 if !fetch_context.is_empty() {
350 context_chunks.push("The following fetched results are available:\n");
351 for context in &fetch_context {
352 context_chunks.push(&context.name);
353 for chunk in &context.text {
354 context_chunks.push(&chunk);
355 }
356 }
357 }
358
359 if !thread_context.is_empty() {
360 context_chunks.push("The following previous conversation threads are available:\n");
361 for context in &thread_context {
362 context_chunks.push(&context.name);
363 for chunk in &context.text {
364 context_chunks.push(&chunk);
365 }
366 }
367 }
368
369 debug_assert!(
370 context_chunks.len() == capacity,
371 "attach_context_message calculated capacity of {}, but length was {}",
372 capacity,
373 context_chunks.len()
374 );
375
376 if !context_chunks.is_empty() {
377 message
378 .content
379 .push(MessageContent::Text(context_chunks.join("\n")));
380 }
381}