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