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) => {
108 let path = file_context.path(cx)?;
109 let full_path: SharedString = path.to_string_lossy().into_owned().into();
110 let name = match path.file_name() {
111 Some(name) => name.to_string_lossy().into_owned().into(),
112 None => full_path.clone(),
113 };
114 let parent = path
115 .parent()
116 .and_then(|p| p.file_name())
117 .map(|p| p.to_string_lossy().into_owned().into());
118
119 Some(ContextSnapshot {
120 id: self.id(),
121 name,
122 parent,
123 tooltip: Some(full_path),
124 kind: ContextKind::File,
125 text: file_context.text.clone(),
126 })
127 }
128 Self::Directory(DirectoryContext { snapshot, .. }) => Some(snapshot.clone()),
129 Self::FetchedUrl(FetchedUrlContext { url, text, id }) => Some(ContextSnapshot {
130 id: *id,
131 name: url.clone(),
132 parent: None,
133 tooltip: None,
134 kind: ContextKind::FetchedUrl,
135 text: text.clone(),
136 }),
137 Self::Thread(thread_context) => {
138 let thread = thread_context.thread.read(cx);
139
140 Some(ContextSnapshot {
141 id: self.id(),
142 name: thread.summary().unwrap_or("New thread".into()),
143 parent: None,
144 tooltip: None,
145 kind: ContextKind::Thread,
146 text: thread_context.text.clone(),
147 })
148 }
149 }
150 }
151}
152
153impl FileContext {
154 pub fn path(&self, cx: &AppContext) -> Option<Arc<Path>> {
155 let buffer = self.buffer.read(cx);
156 if let Some(file) = buffer.file() {
157 Some(file.path().clone())
158 } else {
159 log::error!("Buffer that had a path unexpectedly no longer has a path.");
160 None
161 }
162 }
163}
164
165pub fn attach_context_to_message(
166 message: &mut LanguageModelRequestMessage,
167 contexts: impl Iterator<Item = ContextSnapshot>,
168) {
169 let mut file_context = String::new();
170 let mut directory_context = String::new();
171 let mut fetch_context = String::new();
172 let mut thread_context = String::new();
173
174 for context in contexts {
175 match context.kind {
176 ContextKind::File => {
177 file_context.push_str(&context.text);
178 file_context.push('\n');
179 }
180 ContextKind::Directory => {
181 directory_context.push_str(&context.text);
182 directory_context.push('\n');
183 }
184 ContextKind::FetchedUrl => {
185 fetch_context.push_str(&context.name);
186 fetch_context.push('\n');
187 fetch_context.push_str(&context.text);
188 fetch_context.push('\n');
189 }
190 ContextKind::Thread { .. } => {
191 thread_context.push_str(&context.name);
192 thread_context.push('\n');
193 thread_context.push_str(&context.text);
194 thread_context.push('\n');
195 }
196 }
197 }
198
199 let mut context_text = String::new();
200 if !file_context.is_empty() {
201 context_text.push_str("The following files are available:\n");
202 context_text.push_str(&file_context);
203 }
204
205 if !directory_context.is_empty() {
206 context_text.push_str("The following directories are available:\n");
207 context_text.push_str(&directory_context);
208 }
209
210 if !fetch_context.is_empty() {
211 context_text.push_str("The following fetched results are available\n");
212 context_text.push_str(&fetch_context);
213 }
214
215 if !thread_context.is_empty() {
216 context_text.push_str("The following previous conversation threads are available\n");
217 context_text.push_str(&thread_context);
218 }
219
220 if !context_text.is_empty() {
221 message.content.push(MessageContent::Text(context_text));
222 }
223}