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