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