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