1use std::{ops::Range, path::Path, sync::Arc};
  2
  3use gpui::{App, Entity, SharedString};
  4use language::{Buffer, File};
  5use language_model::LanguageModelRequestMessage;
  6use project::{ProjectEntryId, ProjectPath, Worktree};
  7use prompt_store::UserPromptId;
  8use rope::Point;
  9use serde::{Deserialize, Serialize};
 10use text::{Anchor, BufferId};
 11use ui::IconName;
 12use util::post_inc;
 13
 14use crate::thread::Thread;
 15
 16pub const RULES_ICON: IconName = IconName::Context;
 17
 18#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
 19pub struct ContextId(pub(crate) usize);
 20
 21impl ContextId {
 22    pub fn post_inc(&mut self) -> Self {
 23        Self(post_inc(&mut self.0))
 24    }
 25}
 26
 27pub enum ContextKind {
 28    File,
 29    Directory,
 30    Symbol,
 31    Excerpt,
 32    FetchedUrl,
 33    Thread,
 34    Rules,
 35}
 36
 37impl ContextKind {
 38    pub fn icon(&self) -> IconName {
 39        match self {
 40            ContextKind::File => IconName::File,
 41            ContextKind::Directory => IconName::Folder,
 42            ContextKind::Symbol => IconName::Code,
 43            ContextKind::Excerpt => IconName::Code,
 44            ContextKind::FetchedUrl => IconName::Globe,
 45            ContextKind::Thread => IconName::MessageBubbles,
 46            ContextKind::Rules => RULES_ICON,
 47        }
 48    }
 49}
 50
 51#[derive(Debug, Clone)]
 52pub enum AssistantContext {
 53    File(FileContext),
 54    Directory(DirectoryContext),
 55    Symbol(SymbolContext),
 56    FetchedUrl(FetchedUrlContext),
 57    Thread(ThreadContext),
 58    Excerpt(ExcerptContext),
 59    Rules(RulesContext),
 60}
 61
 62impl AssistantContext {
 63    pub fn id(&self) -> ContextId {
 64        match self {
 65            Self::File(file) => file.id,
 66            Self::Directory(directory) => directory.id,
 67            Self::Symbol(symbol) => symbol.id,
 68            Self::FetchedUrl(url) => url.id,
 69            Self::Thread(thread) => thread.id,
 70            Self::Excerpt(excerpt) => excerpt.id,
 71            Self::Rules(rules) => rules.id,
 72        }
 73    }
 74}
 75
 76#[derive(Debug, Clone)]
 77pub struct FileContext {
 78    pub id: ContextId,
 79    pub context_buffer: ContextBuffer,
 80}
 81
 82#[derive(Debug, Clone)]
 83pub struct DirectoryContext {
 84    pub id: ContextId,
 85    pub worktree: Entity<Worktree>,
 86    pub entry_id: ProjectEntryId,
 87    pub last_path: Arc<Path>,
 88    /// Buffers of the files within the directory.
 89    pub context_buffers: Vec<ContextBuffer>,
 90}
 91
 92impl DirectoryContext {
 93    pub fn entry<'a>(&self, cx: &'a App) -> Option<&'a project::Entry> {
 94        self.worktree.read(cx).entry_for_id(self.entry_id)
 95    }
 96
 97    pub fn project_path(&self, cx: &App) -> Option<ProjectPath> {
 98        let worktree = self.worktree.read(cx);
 99        worktree
100            .entry_for_id(self.entry_id)
101            .map(|entry| ProjectPath {
102                worktree_id: worktree.id(),
103                path: entry.path.clone(),
104            })
105    }
106}
107
108#[derive(Debug, Clone)]
109pub struct SymbolContext {
110    pub id: ContextId,
111    pub context_symbol: ContextSymbol,
112}
113
114#[derive(Debug, Clone)]
115pub struct FetchedUrlContext {
116    pub id: ContextId,
117    pub url: SharedString,
118    pub text: SharedString,
119}
120
121#[derive(Debug, Clone)]
122pub struct ThreadContext {
123    pub id: ContextId,
124    // TODO: Entity<Thread> holds onto the thread even if the thread is deleted. Should probably be
125    // a WeakEntity and handle removal from the UI when it has dropped.
126    pub thread: Entity<Thread>,
127    pub text: SharedString,
128}
129
130impl ThreadContext {
131    pub fn summary(&self, cx: &App) -> SharedString {
132        self.thread
133            .read(cx)
134            .summary()
135            .unwrap_or("New thread".into())
136    }
137}
138
139#[derive(Clone)]
140pub struct ContextBuffer {
141    pub id: BufferId,
142    // TODO: Entity<Buffer> holds onto the buffer even if the buffer is deleted. Should probably be
143    // a WeakEntity and handle removal from the UI when it has dropped.
144    pub buffer: Entity<Buffer>,
145    pub file: Arc<dyn File>,
146    pub version: clock::Global,
147    pub text: SharedString,
148}
149
150impl std::fmt::Debug for ContextBuffer {
151    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
152        f.debug_struct("ContextBuffer")
153            .field("id", &self.id)
154            .field("buffer", &self.buffer)
155            .field("version", &self.version)
156            .field("text", &self.text)
157            .finish()
158    }
159}
160
161#[derive(Debug, Clone)]
162pub struct ContextSymbol {
163    pub id: ContextSymbolId,
164    pub buffer: Entity<Buffer>,
165    pub buffer_version: clock::Global,
166    /// The range that the symbol encloses, e.g. for function symbol, this will
167    /// include not only the signature, but also the body
168    pub enclosing_range: Range<Anchor>,
169    pub text: SharedString,
170}
171
172#[derive(Debug, Clone, PartialEq, Eq, Hash)]
173pub struct ContextSymbolId {
174    pub path: ProjectPath,
175    pub name: SharedString,
176    pub range: Range<Anchor>,
177}
178
179#[derive(Debug, Clone)]
180pub struct ExcerptContext {
181    pub id: ContextId,
182    pub range: Range<Anchor>,
183    pub line_range: Range<Point>,
184    pub context_buffer: ContextBuffer,
185}
186
187#[derive(Debug, Clone)]
188pub struct RulesContext {
189    pub id: ContextId,
190    pub prompt_id: UserPromptId,
191    pub title: SharedString,
192    pub text: SharedString,
193}
194
195/// Formats a collection of contexts into a string representation
196pub fn format_context_as_string<'a>(
197    contexts: impl Iterator<Item = &'a AssistantContext>,
198    cx: &App,
199) -> Option<String> {
200    let mut file_context = Vec::new();
201    let mut directory_context = Vec::new();
202    let mut symbol_context = Vec::new();
203    let mut excerpt_context = Vec::new();
204    let mut fetch_context = Vec::new();
205    let mut thread_context = Vec::new();
206    let mut rules_context = Vec::new();
207
208    for context in contexts {
209        match context {
210            AssistantContext::File(context) => file_context.push(context),
211            AssistantContext::Directory(context) => directory_context.push(context),
212            AssistantContext::Symbol(context) => symbol_context.push(context),
213            AssistantContext::Excerpt(context) => excerpt_context.push(context),
214            AssistantContext::FetchedUrl(context) => fetch_context.push(context),
215            AssistantContext::Thread(context) => thread_context.push(context),
216            AssistantContext::Rules(context) => rules_context.push(context),
217        }
218    }
219
220    if file_context.is_empty()
221        && directory_context.is_empty()
222        && symbol_context.is_empty()
223        && excerpt_context.is_empty()
224        && fetch_context.is_empty()
225        && thread_context.is_empty()
226        && rules_context.is_empty()
227    {
228        return None;
229    }
230
231    let mut result = String::new();
232    result.push_str("\n<context>\n\
233        The following items were attached by the user. You don't need to use other tools to read them.\n\n");
234
235    if !file_context.is_empty() {
236        result.push_str("<files>\n");
237        for context in file_context {
238            result.push_str(&context.context_buffer.text);
239        }
240        result.push_str("</files>\n");
241    }
242
243    if !directory_context.is_empty() {
244        result.push_str("<directories>\n");
245        for context in directory_context {
246            for context_buffer in &context.context_buffers {
247                result.push_str(&context_buffer.text);
248            }
249        }
250        result.push_str("</directories>\n");
251    }
252
253    if !symbol_context.is_empty() {
254        result.push_str("<symbols>\n");
255        for context in symbol_context {
256            result.push_str(&context.context_symbol.text);
257            result.push('\n');
258        }
259        result.push_str("</symbols>\n");
260    }
261
262    if !excerpt_context.is_empty() {
263        result.push_str("<excerpts>\n");
264        for context in excerpt_context {
265            result.push_str(&context.context_buffer.text);
266            result.push('\n');
267        }
268        result.push_str("</excerpts>\n");
269    }
270
271    if !fetch_context.is_empty() {
272        result.push_str("<fetched_urls>\n");
273        for context in &fetch_context {
274            result.push_str(&context.url);
275            result.push('\n');
276            result.push_str(&context.text);
277            result.push('\n');
278        }
279        result.push_str("</fetched_urls>\n");
280    }
281
282    if !thread_context.is_empty() {
283        result.push_str("<conversation_threads>\n");
284        for context in &thread_context {
285            result.push_str(&context.summary(cx));
286            result.push('\n');
287            result.push_str(&context.text);
288            result.push('\n');
289        }
290        result.push_str("</conversation_threads>\n");
291    }
292
293    if !rules_context.is_empty() {
294        result.push_str(
295            "<user_rules>\n\
296            The user has specified the following rules that should be applied:\n\n",
297        );
298        for context in &rules_context {
299            result.push_str(&context.text);
300            result.push('\n');
301        }
302        result.push_str("</user_rules>\n");
303    }
304
305    result.push_str("</context>\n");
306    Some(result)
307}
308
309pub fn attach_context_to_message<'a>(
310    message: &mut LanguageModelRequestMessage,
311    contexts: impl Iterator<Item = &'a AssistantContext>,
312    cx: &App,
313) {
314    if let Some(context_string) = format_context_as_string(contexts, cx) {
315        message.content.push(context_string.into());
316    }
317}