1use std::path::Path;
2use std::rc::Rc;
3use std::sync::Arc;
4
5use file_icons::FileIcons;
6use gpui::{AppContext, Model, SharedString};
7use language::Buffer;
8use language_model::{LanguageModelRequestMessage, MessageContent};
9use serde::{Deserialize, Serialize};
10use text::BufferId;
11use ui::IconName;
12use util::post_inc;
13
14use crate::thread::Thread;
15
16#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
17pub struct ContextId(pub(crate) usize);
18
19impl ContextId {
20 pub fn post_inc(&mut self) -> Self {
21 Self(post_inc(&mut self.0))
22 }
23}
24
25/// Some context attached to a message in a thread.
26#[derive(Debug, Clone)]
27pub struct ContextSnapshot {
28 pub id: ContextId,
29 pub name: SharedString,
30 pub parent: Option<SharedString>,
31 pub tooltip: Option<SharedString>,
32 pub icon_path: Option<SharedString>,
33 pub kind: ContextKind,
34 /// Joining these strings separated by \n yields text for model. Not refreshed by `snapshot`.
35 pub text: Box<[SharedString]>,
36}
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum ContextKind {
40 File,
41 Directory,
42 FetchedUrl,
43 Thread,
44}
45
46impl ContextKind {
47 pub fn icon(&self) -> IconName {
48 match self {
49 ContextKind::File => IconName::File,
50 ContextKind::Directory => IconName::Folder,
51 ContextKind::FetchedUrl => IconName::Globe,
52 ContextKind::Thread => IconName::MessageCircle,
53 }
54 }
55}
56
57#[derive(Debug)]
58pub enum Context {
59 File(FileContext),
60 Directory(DirectoryContext),
61 FetchedUrl(FetchedUrlContext),
62 Thread(ThreadContext),
63}
64
65impl Context {
66 pub fn id(&self) -> ContextId {
67 match self {
68 Self::File(file) => file.id,
69 Self::Directory(directory) => directory.snapshot.id,
70 Self::FetchedUrl(url) => url.id,
71 Self::Thread(thread) => thread.id,
72 }
73 }
74}
75
76#[derive(Debug)]
77pub struct FileContext {
78 pub id: ContextId,
79 pub buffer: ContextBuffer,
80}
81
82#[derive(Debug)]
83pub struct DirectoryContext {
84 #[allow(unused)]
85 pub path: Rc<Path>,
86 #[allow(unused)]
87 pub buffers: Vec<ContextBuffer>,
88 pub snapshot: ContextSnapshot,
89}
90
91#[derive(Debug)]
92pub struct FetchedUrlContext {
93 pub id: ContextId,
94 pub url: SharedString,
95 pub text: SharedString,
96}
97
98// TODO: Model<Thread> holds onto the thread even if the thread is deleted. Can either handle this
99// explicitly or have a WeakModel<Thread> and remove during snapshot.
100
101#[derive(Debug)]
102pub struct ThreadContext {
103 pub id: ContextId,
104 pub thread: Model<Thread>,
105 pub text: SharedString,
106}
107
108// TODO: Model<Buffer> holds onto the buffer even if the file is deleted and closed. Should remove
109// the context from the message editor in this case.
110
111#[derive(Debug)]
112pub struct ContextBuffer {
113 #[allow(unused)]
114 pub id: BufferId,
115 pub buffer: Model<Buffer>,
116 #[allow(unused)]
117 pub version: clock::Global,
118 pub text: SharedString,
119}
120
121impl Context {
122 pub fn snapshot(&self, cx: &AppContext) -> Option<ContextSnapshot> {
123 match &self {
124 Self::File(file_context) => file_context.snapshot(cx),
125 Self::Directory(directory_context) => Some(directory_context.snapshot()),
126 Self::FetchedUrl(fetched_url_context) => Some(fetched_url_context.snapshot()),
127 Self::Thread(thread_context) => Some(thread_context.snapshot(cx)),
128 }
129 }
130}
131
132impl FileContext {
133 pub fn path(&self, cx: &AppContext) -> Option<Arc<Path>> {
134 let buffer = self.buffer.buffer.read(cx);
135 if let Some(file) = buffer.file() {
136 Some(file.path().clone())
137 } else {
138 log::error!("Buffer that had a path unexpectedly no longer has a path.");
139 None
140 }
141 }
142
143 pub fn snapshot(&self, cx: &AppContext) -> Option<ContextSnapshot> {
144 let path = self.path(cx)?;
145 let full_path: SharedString = path.to_string_lossy().into_owned().into();
146 let name = match path.file_name() {
147 Some(name) => name.to_string_lossy().into_owned().into(),
148 None => full_path.clone(),
149 };
150 let parent = path
151 .parent()
152 .and_then(|p| p.file_name())
153 .map(|p| p.to_string_lossy().into_owned().into());
154
155 let icon_path = FileIcons::get_icon(&path, cx);
156
157 Some(ContextSnapshot {
158 id: self.id,
159 name,
160 parent,
161 tooltip: Some(full_path),
162 icon_path,
163 kind: ContextKind::File,
164 text: Box::new([self.buffer.text.clone()]),
165 })
166 }
167}
168
169impl DirectoryContext {
170 pub fn snapshot(&self) -> ContextSnapshot {
171 self.snapshot.clone()
172 }
173}
174
175impl FetchedUrlContext {
176 pub fn snapshot(&self) -> ContextSnapshot {
177 ContextSnapshot {
178 id: self.id,
179 name: self.url.clone(),
180 parent: None,
181 tooltip: None,
182 icon_path: None,
183 kind: ContextKind::FetchedUrl,
184 text: Box::new([self.text.clone()]),
185 }
186 }
187}
188
189impl ThreadContext {
190 pub fn snapshot(&self, cx: &AppContext) -> ContextSnapshot {
191 let thread = self.thread.read(cx);
192 ContextSnapshot {
193 id: self.id,
194 name: thread.summary().unwrap_or("New thread".into()),
195 parent: None,
196 tooltip: None,
197 icon_path: None,
198 kind: ContextKind::Thread,
199 text: Box::new([self.text.clone()]),
200 }
201 }
202}
203
204pub fn attach_context_to_message(
205 message: &mut LanguageModelRequestMessage,
206 contexts: impl Iterator<Item = ContextSnapshot>,
207) {
208 let mut file_context = Vec::new();
209 let mut directory_context = Vec::new();
210 let mut fetch_context = Vec::new();
211 let mut thread_context = Vec::new();
212
213 let mut capacity = 0;
214 for context in contexts {
215 capacity += context.text.len();
216 match context.kind {
217 ContextKind::File => file_context.push(context),
218 ContextKind::Directory => directory_context.push(context),
219 ContextKind::FetchedUrl => fetch_context.push(context),
220 ContextKind::Thread => thread_context.push(context),
221 }
222 }
223 if !file_context.is_empty() {
224 capacity += 1;
225 }
226 if !directory_context.is_empty() {
227 capacity += 1;
228 }
229 if !fetch_context.is_empty() {
230 capacity += 1 + fetch_context.len();
231 }
232 if !thread_context.is_empty() {
233 capacity += 1 + thread_context.len();
234 }
235 if capacity == 0 {
236 return;
237 }
238
239 let mut context_chunks = Vec::with_capacity(capacity);
240
241 if !file_context.is_empty() {
242 context_chunks.push("The following files are available:\n");
243 for context in &file_context {
244 for chunk in &context.text {
245 context_chunks.push(&chunk);
246 }
247 }
248 }
249
250 if !directory_context.is_empty() {
251 context_chunks.push("The following directories are available:\n");
252 for context in &directory_context {
253 for chunk in &context.text {
254 context_chunks.push(&chunk);
255 }
256 }
257 }
258
259 if !fetch_context.is_empty() {
260 context_chunks.push("The following fetched results are available:\n");
261 for context in &fetch_context {
262 context_chunks.push(&context.name);
263 for chunk in &context.text {
264 context_chunks.push(&chunk);
265 }
266 }
267 }
268
269 if !thread_context.is_empty() {
270 context_chunks.push("The following previous conversation threads are available:\n");
271 for context in &thread_context {
272 context_chunks.push(&context.name);
273 for chunk in &context.text {
274 context_chunks.push(&chunk);
275 }
276 }
277 }
278
279 debug_assert!(
280 context_chunks.len() == capacity,
281 "attach_context_message calculated capacity of {}, but length was {}",
282 capacity,
283 context_chunks.len()
284 );
285
286 if !context_chunks.is_empty() {
287 message
288 .content
289 .push(MessageContent::Text(context_chunks.join("\n")));
290 }
291}