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