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