1use std::path::Path;
2use std::rc::Rc;
3
4use file_icons::FileIcons;
5use gpui::{App, Entity, SharedString};
6use language::Buffer;
7use language_model::{LanguageModelRequestMessage, MessageContent};
8use serde::{Deserialize, Serialize};
9use text::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 FetchedUrl,
42 Thread,
43}
44
45impl ContextKind {
46 pub fn label(&self) -> &'static str {
47 match self {
48 ContextKind::File => "File",
49 ContextKind::Directory => "Folder",
50 ContextKind::FetchedUrl => "Fetch",
51 ContextKind::Thread => "Thread",
52 }
53 }
54
55 pub fn icon(&self) -> IconName {
56 match self {
57 ContextKind::File => IconName::File,
58 ContextKind::Directory => IconName::Folder,
59 ContextKind::FetchedUrl => IconName::Globe,
60 ContextKind::Thread => IconName::MessageCircle,
61 }
62 }
63}
64
65#[derive(Debug)]
66pub enum AssistantContext {
67 File(FileContext),
68 Directory(DirectoryContext),
69 FetchedUrl(FetchedUrlContext),
70 Thread(ThreadContext),
71}
72
73impl AssistantContext {
74 pub fn id(&self) -> ContextId {
75 match self {
76 Self::File(file) => file.id,
77 Self::Directory(directory) => directory.snapshot.id,
78 Self::FetchedUrl(url) => url.id,
79 Self::Thread(thread) => thread.id,
80 }
81 }
82}
83
84#[derive(Debug)]
85pub struct FileContext {
86 pub id: ContextId,
87 pub context_buffer: ContextBuffer,
88}
89
90#[derive(Debug)]
91pub struct DirectoryContext {
92 pub path: Rc<Path>,
93 pub context_buffers: Vec<ContextBuffer>,
94 pub snapshot: ContextSnapshot,
95}
96
97#[derive(Debug)]
98pub struct FetchedUrlContext {
99 pub id: ContextId,
100 pub url: SharedString,
101 pub text: SharedString,
102}
103
104// TODO: Model<Thread> holds onto the thread even if the thread is deleted. Can either handle this
105// explicitly or have a WeakModel<Thread> and remove during snapshot.
106
107#[derive(Debug)]
108pub struct ThreadContext {
109 pub id: ContextId,
110 pub thread: Entity<Thread>,
111 pub text: SharedString,
112}
113
114// TODO: Model<Buffer> holds onto the buffer even if the file is deleted and closed. Should remove
115// the context from the message editor in this case.
116
117#[derive(Debug, Clone)]
118pub struct ContextBuffer {
119 pub id: BufferId,
120 pub buffer: Entity<Buffer>,
121 pub version: clock::Global,
122 pub text: SharedString,
123}
124
125impl AssistantContext {
126 pub fn snapshot(&self, cx: &App) -> Option<ContextSnapshot> {
127 match &self {
128 Self::File(file_context) => file_context.snapshot(cx),
129 Self::Directory(directory_context) => Some(directory_context.snapshot()),
130 Self::FetchedUrl(fetched_url_context) => Some(fetched_url_context.snapshot()),
131 Self::Thread(thread_context) => Some(thread_context.snapshot(cx)),
132 }
133 }
134}
135
136impl FileContext {
137 pub fn snapshot(&self, cx: &App) -> Option<ContextSnapshot> {
138 let buffer = self.context_buffer.buffer.read(cx);
139 let path = buffer_path_log_err(buffer)?;
140 let full_path: SharedString = path.to_string_lossy().into_owned().into();
141 let name = match path.file_name() {
142 Some(name) => name.to_string_lossy().into_owned().into(),
143 None => full_path.clone(),
144 };
145 let parent = path
146 .parent()
147 .and_then(|p| p.file_name())
148 .map(|p| p.to_string_lossy().into_owned().into());
149
150 let icon_path = FileIcons::get_icon(&path, cx);
151
152 Some(ContextSnapshot {
153 id: self.id,
154 name,
155 parent,
156 tooltip: Some(full_path),
157 icon_path,
158 kind: ContextKind::File,
159 text: Box::new([self.context_buffer.text.clone()]),
160 })
161 }
162}
163
164impl DirectoryContext {
165 pub fn new(
166 id: ContextId,
167 path: &Path,
168 context_buffers: Vec<ContextBuffer>,
169 ) -> DirectoryContext {
170 let full_path: SharedString = path.to_string_lossy().into_owned().into();
171
172 let name = match path.file_name() {
173 Some(name) => name.to_string_lossy().into_owned().into(),
174 None => full_path.clone(),
175 };
176
177 let parent = path
178 .parent()
179 .and_then(|p| p.file_name())
180 .map(|p| p.to_string_lossy().into_owned().into());
181
182 // TODO: include directory path in text?
183 let text = context_buffers
184 .iter()
185 .map(|b| b.text.clone())
186 .collect::<Vec<_>>()
187 .into();
188
189 DirectoryContext {
190 path: path.into(),
191 context_buffers,
192 snapshot: ContextSnapshot {
193 id,
194 name,
195 parent,
196 tooltip: Some(full_path),
197 icon_path: None,
198 kind: ContextKind::Directory,
199 text,
200 },
201 }
202 }
203
204 pub fn snapshot(&self) -> ContextSnapshot {
205 self.snapshot.clone()
206 }
207}
208
209impl FetchedUrlContext {
210 pub fn snapshot(&self) -> ContextSnapshot {
211 ContextSnapshot {
212 id: self.id,
213 name: self.url.clone(),
214 parent: None,
215 tooltip: None,
216 icon_path: None,
217 kind: ContextKind::FetchedUrl,
218 text: Box::new([self.text.clone()]),
219 }
220 }
221}
222
223impl ThreadContext {
224 pub fn snapshot(&self, cx: &App) -> ContextSnapshot {
225 let thread = self.thread.read(cx);
226 ContextSnapshot {
227 id: self.id,
228 name: thread.summary().unwrap_or("New thread".into()),
229 parent: None,
230 tooltip: None,
231 icon_path: None,
232 kind: ContextKind::Thread,
233 text: Box::new([self.text.clone()]),
234 }
235 }
236}
237
238pub fn attach_context_to_message(
239 message: &mut LanguageModelRequestMessage,
240 contexts: impl Iterator<Item = ContextSnapshot>,
241) {
242 let mut file_context = Vec::new();
243 let mut directory_context = Vec::new();
244 let mut fetch_context = Vec::new();
245 let mut thread_context = Vec::new();
246
247 let mut capacity = 0;
248 for context in contexts {
249 capacity += context.text.len();
250 match context.kind {
251 ContextKind::File => file_context.push(context),
252 ContextKind::Directory => directory_context.push(context),
253 ContextKind::FetchedUrl => fetch_context.push(context),
254 ContextKind::Thread => thread_context.push(context),
255 }
256 }
257 if !file_context.is_empty() {
258 capacity += 1;
259 }
260 if !directory_context.is_empty() {
261 capacity += 1;
262 }
263 if !fetch_context.is_empty() {
264 capacity += 1 + fetch_context.len();
265 }
266 if !thread_context.is_empty() {
267 capacity += 1 + thread_context.len();
268 }
269 if capacity == 0 {
270 return;
271 }
272
273 let mut context_chunks = Vec::with_capacity(capacity);
274
275 if !file_context.is_empty() {
276 context_chunks.push("The following files are available:\n");
277 for context in &file_context {
278 for chunk in &context.text {
279 context_chunks.push(&chunk);
280 }
281 }
282 }
283
284 if !directory_context.is_empty() {
285 context_chunks.push("The following directories are available:\n");
286 for context in &directory_context {
287 for chunk in &context.text {
288 context_chunks.push(&chunk);
289 }
290 }
291 }
292
293 if !fetch_context.is_empty() {
294 context_chunks.push("The following fetched results are available:\n");
295 for context in &fetch_context {
296 context_chunks.push(&context.name);
297 for chunk in &context.text {
298 context_chunks.push(&chunk);
299 }
300 }
301 }
302
303 if !thread_context.is_empty() {
304 context_chunks.push("The following previous conversation threads are available:\n");
305 for context in &thread_context {
306 context_chunks.push(&context.name);
307 for chunk in &context.text {
308 context_chunks.push(&chunk);
309 }
310 }
311 }
312
313 debug_assert!(
314 context_chunks.len() == capacity,
315 "attach_context_message calculated capacity of {}, but length was {}",
316 capacity,
317 context_chunks.len()
318 );
319
320 if !context_chunks.is_empty() {
321 message
322 .content
323 .push(MessageContent::Text(context_chunks.join("\n")));
324 }
325}