use std::{ ops::Range, path::{Path, PathBuf}, sync::Arc,}; use futures::{FutureExt, future::Shared};use gpui::{App, Entity, SharedString, Task};use language::Buffer;use language_model::{LanguageModelImage, LanguageModelRequestMessage};use project::{ProjectEntryId, ProjectPath, Worktree};use prompt_store::UserPromptId;use rope::Point;use serde::{Deserialize, Serialize};use text::{Anchor, BufferId};use ui::IconName;use util::post_inc; use crate::thread::Thread; pub const RULES_ICON: IconName = IconName::Context; #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]pub struct ContextId(pub(crate) usize); impl ContextId { pub fn post_inc(&mut self) -> Self { Self(post_inc(&mut self.0)) }} pub enum ContextKind { File, Directory, Symbol, Selection, FetchedUrl, Thread, Rules, Image,} impl ContextKind { pub fn icon(&self) -> IconName { match self { ContextKind::File => IconName::File, ContextKind::Directory => IconName::Folder, ContextKind::Symbol => IconName::Code, ContextKind::Selection => IconName::Context, ContextKind::FetchedUrl => IconName::Globe, ContextKind::Thread => IconName::MessageBubbles, ContextKind::Rules => RULES_ICON, ContextKind::Image => IconName::Image, } }} #[derive(Debug, Clone)]pub enum AssistantContext { File(FileContext), Directory(DirectoryContext), Symbol(SymbolContext), FetchedUrl(FetchedUrlContext), Thread(ThreadContext), Selection(SelectionContext), Rules(RulesContext), Image(ImageContext),} impl AssistantContext { pub fn id(&self) -> ContextId { match self { Self::File(file) => file.id, Self::Directory(directory) => directory.id, Self::Symbol(symbol) => symbol.id, Self::FetchedUrl(url) => url.id, Self::Thread(thread) => thread.id, Self::Selection(selection) => selection.id, Self::Rules(rules) => rules.id, Self::Image(image) => image.id, } }} #[derive(Debug, Clone)]pub struct FileContext { pub id: ContextId, pub context_buffer: ContextBuffer,} #[derive(Debug, Clone)]pub struct DirectoryContext { pub id: ContextId, pub worktree: Entity<Worktree>, pub entry_id: ProjectEntryId, pub last_path: Arc<Path>, /// Buffers of the files within the directory. pub context_buffers: Vec<ContextBuffer>,} impl DirectoryContext { pub fn entry<'a>(&self, cx: &'a App) -> Option<&'a project::Entry> { self.worktree.read(cx).entry_for_id(self.entry_id) } pub fn project_path(&self, cx: &App) -> Option<ProjectPath> { let worktree = self.worktree.read(cx); worktree .entry_for_id(self.entry_id) .map(|entry| ProjectPath { worktree_id: worktree.id(), path: entry.path.clone(), }) }} #[derive(Debug, Clone)]pub struct SymbolContext { pub id: ContextId, pub context_symbol: ContextSymbol,} #[derive(Debug, Clone)]pub struct FetchedUrlContext { pub id: ContextId, pub url: SharedString, pub text: SharedString,} #[derive(Debug, Clone)]pub struct ThreadContext { pub id: ContextId, // TODO: Entity<Thread> holds onto the thread even if the thread is deleted. Should probably be // a WeakEntity and handle removal from the UI when it has dropped. pub thread: Entity<Thread>, pub text: SharedString,} impl ThreadContext { pub fn summary(&self, cx: &App) -> SharedString { self.thread .read(cx) .summary() .unwrap_or("New thread".into()) }} #[derive(Debug, Clone)]pub struct ImageContext { pub id: ContextId, pub original_image: Arc<gpui::Image>, pub image_task: Shared<Task<Option<LanguageModelImage>>>,} impl ImageContext { pub fn image(&self) -> Option<LanguageModelImage> { self.image_task.clone().now_or_never().flatten() } pub fn is_loading(&self) -> bool { self.image_task.clone().now_or_never().is_none() } pub fn is_error(&self) -> bool { self.image_task .clone() .now_or_never() .map(|result| result.is_none()) .unwrap_or(false) }} #[derive(Clone)]pub struct ContextBuffer { pub id: BufferId, // TODO: Entity<Buffer> holds onto the buffer even if the buffer is deleted. Should probably be // a WeakEntity and handle removal from the UI when it has dropped. pub buffer: Entity<Buffer>, pub last_full_path: Arc<Path>, pub version: clock::Global, pub text: SharedString,} impl ContextBuffer { pub fn full_path(&self, cx: &App) -> PathBuf { let file = self.buffer.read(cx).file(); // Note that in practice file can't be `None` because it is present when this is created and // there's no way for buffers to go from having a file to not. file.map_or(self.last_full_path.to_path_buf(), |file| file.full_path(cx)) }} impl std::fmt::Debug for ContextBuffer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ContextBuffer") .field("id", &self.id) .field("buffer", &self.buffer) .field("version", &self.version) .field("text", &self.text) .finish() }} #[derive(Debug, Clone)]pub struct ContextSymbol { pub id: ContextSymbolId, pub buffer: Entity<Buffer>, pub buffer_version: clock::Global, /// The range that the symbol encloses, e.g. for function symbol, this will /// include not only the signature, but also the body pub enclosing_range: Range<Anchor>, pub text: SharedString,} #[derive(Debug, Clone, PartialEq, Eq, Hash)]pub struct ContextSymbolId { pub path: ProjectPath, pub name: SharedString, pub range: Range<Anchor>,} #[derive(Debug, Clone)]pub struct SelectionContext { pub id: ContextId, pub range: Range<Anchor>, pub line_range: Range<Point>, pub context_buffer: ContextBuffer,} #[derive(Debug, Clone)]pub struct RulesContext { pub id: ContextId, pub prompt_id: UserPromptId, pub title: SharedString, pub text: SharedString,} /// Formats a collection of contexts into a string representationpub fn format_context_as_string<'a>( contexts: impl Iterator<Item = &'a AssistantContext>, cx: &App,) -> Option<String> { let mut file_context = Vec::new(); let mut directory_context = Vec::new(); let mut symbol_context = Vec::new(); let mut selection_context = Vec::new(); let mut fetch_context = Vec::new(); let mut thread_context = Vec::new(); let mut rules_context = Vec::new(); for context in contexts { match context { AssistantContext::File(context) => file_context.push(context), AssistantContext::Directory(context) => directory_context.push(context), AssistantContext::Symbol(context) => symbol_context.push(context), AssistantContext::Selection(context) => selection_context.push(context), AssistantContext::FetchedUrl(context) => fetch_context.push(context), AssistantContext::Thread(context) => thread_context.push(context), AssistantContext::Rules(context) => rules_context.push(context), AssistantContext::Image(_) => {} } } if file_context.is_empty() && directory_context.is_empty() && symbol_context.is_empty() && selection_context.is_empty() && fetch_context.is_empty() && thread_context.is_empty() && rules_context.is_empty() { return None; } let mut result = String::new(); result.push_str("\n<context>\n\ The following items were attached by the user. You don't need to use other tools to read them.\n\n"); if !file_context.is_empty() { result.push_str("<files>\n"); for context in file_context { result.push_str(&context.context_buffer.text); } result.push_str("</files>\n"); } if !directory_context.is_empty() { result.push_str("<directories>\n"); for context in directory_context { for context_buffer in &context.context_buffers { result.push_str(&context_buffer.text); } } result.push_str("</directories>\n"); } if !symbol_context.is_empty() { result.push_str("<symbols>\n"); for context in symbol_context { result.push_str(&context.context_symbol.text); result.push('\n'); } result.push_str("</symbols>\n"); } if !selection_context.is_empty() { result.push_str("<selections>\n"); for context in selection_context { result.push_str(&context.context_buffer.text); result.push('\n'); } result.push_str("</selections>\n"); } if !fetch_context.is_empty() { result.push_str("<fetched_urls>\n"); for context in &fetch_context { result.push_str(&context.url); result.push('\n'); result.push_str(&context.text); result.push('\n'); } result.push_str("</fetched_urls>\n"); } if !thread_context.is_empty() { result.push_str("<conversation_threads>\n"); for context in &thread_context { result.push_str(&context.summary(cx)); result.push('\n'); result.push_str(&context.text); result.push('\n'); } result.push_str("</conversation_threads>\n"); } if !rules_context.is_empty() { result.push_str( "<user_rules>\n\ The user has specified the following rules that should be applied:\n\n", ); for context in &rules_context { result.push_str(&context.text); result.push('\n'); } result.push_str("</user_rules>\n"); } result.push_str("</context>\n"); Some(result)} pub fn attach_context_to_message<'a>( message: &mut LanguageModelRequestMessage, contexts: impl Iterator<Item = &'a AssistantContext>, cx: &App,) { if let Some(context_string) = format_context_as_string(contexts, cx) { message.content.push(context_string.into()); }}
use std::{
ops::Range,
path::{Path, PathBuf},
sync::Arc,
};
use futures::{FutureExt, future::Shared};
use gpui::{App, Entity, SharedString, Task};
use language::Buffer;
use language_model::{LanguageModelImage, LanguageModelRequestMessage};
use project::{ProjectEntryId, ProjectPath, Worktree};
use prompt_store::UserPromptId;
use rope::Point;
use serde::{Deserialize, Serialize};
use text::{Anchor, BufferId};
use ui::IconName;
use util::post_inc;
use crate::thread::Thread;
pub const RULES_ICON: IconName = IconName::Context;
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
pub struct ContextId(pub(crate) usize);
impl ContextId {
pub fn post_inc(&mut self) -> Self {
Self(post_inc(&mut self.0))
}
}
pub enum ContextKind {
File,
Directory,
Symbol,
Selection,
FetchedUrl,
Thread,
Rules,
Image,
}
impl ContextKind {
pub fn icon(&self) -> IconName {
match self {
ContextKind::File => IconName::File,
ContextKind::Directory => IconName::Folder,
ContextKind::Symbol => IconName::Code,
ContextKind::Selection => IconName::Context,
ContextKind::FetchedUrl => IconName::Globe,
ContextKind::Thread => IconName::MessageBubbles,
ContextKind::Rules => RULES_ICON,
ContextKind::Image => IconName::Image,
}
}
}
#[derive(Debug, Clone)]
pub enum AssistantContext {
File(FileContext),
Directory(DirectoryContext),
Symbol(SymbolContext),
FetchedUrl(FetchedUrlContext),
Thread(ThreadContext),
Selection(SelectionContext),
Rules(RulesContext),
Image(ImageContext),
}
impl AssistantContext {
pub fn id(&self) -> ContextId {
match self {
Self::File(file) => file.id,
Self::Directory(directory) => directory.id,
Self::Symbol(symbol) => symbol.id,
Self::FetchedUrl(url) => url.id,
Self::Thread(thread) => thread.id,
Self::Selection(selection) => selection.id,
Self::Rules(rules) => rules.id,
Self::Image(image) => image.id,
}
}
}
#[derive(Debug, Clone)]
pub struct FileContext {
pub id: ContextId,
pub context_buffer: ContextBuffer,
}
#[derive(Debug, Clone)]
pub struct DirectoryContext {
pub id: ContextId,
pub worktree: Entity<Worktree>,
pub entry_id: ProjectEntryId,
pub last_path: Arc<Path>,
/// Buffers of the files within the directory.
pub context_buffers: Vec<ContextBuffer>,
}
impl DirectoryContext {
pub fn entry<'a>(&self, cx: &'a App) -> Option<&'a project::Entry> {
self.worktree.read(cx).entry_for_id(self.entry_id)
}
pub fn project_path(&self, cx: &App) -> Option<ProjectPath> {
let worktree = self.worktree.read(cx);
worktree
.entry_for_id(self.entry_id)
.map(|entry| ProjectPath {
worktree_id: worktree.id(),
path: entry.path.clone(),
})
}
}
#[derive(Debug, Clone)]
pub struct SymbolContext {
pub id: ContextId,
pub context_symbol: ContextSymbol,
}
#[derive(Debug, Clone)]
pub struct FetchedUrlContext {
pub id: ContextId,
pub url: SharedString,
pub text: SharedString,
}
#[derive(Debug, Clone)]
pub struct ThreadContext {
pub id: ContextId,
// TODO: Entity<Thread> holds onto the thread even if the thread is deleted. Should probably be
// a WeakEntity and handle removal from the UI when it has dropped.
pub thread: Entity<Thread>,
pub text: SharedString,
}
impl ThreadContext {
pub fn summary(&self, cx: &App) -> SharedString {
self.thread
.read(cx)
.summary()
.unwrap_or("New thread".into())
}
}
#[derive(Debug, Clone)]
pub struct ImageContext {
pub id: ContextId,
pub original_image: Arc<gpui::Image>,
pub image_task: Shared<Task<Option<LanguageModelImage>>>,
}
impl ImageContext {
pub fn image(&self) -> Option<LanguageModelImage> {
self.image_task.clone().now_or_never().flatten()
}
pub fn is_loading(&self) -> bool {
self.image_task.clone().now_or_never().is_none()
}
pub fn is_error(&self) -> bool {
self.image_task
.clone()
.now_or_never()
.map(|result| result.is_none())
.unwrap_or(false)
}
}
#[derive(Clone)]
pub struct ContextBuffer {
pub id: BufferId,
// TODO: Entity<Buffer> holds onto the buffer even if the buffer is deleted. Should probably be
// a WeakEntity and handle removal from the UI when it has dropped.
pub buffer: Entity<Buffer>,
pub last_full_path: Arc<Path>,
pub version: clock::Global,
pub text: SharedString,
}
impl ContextBuffer {
pub fn full_path(&self, cx: &App) -> PathBuf {
let file = self.buffer.read(cx).file();
// Note that in practice file can't be `None` because it is present when this is created and
// there's no way for buffers to go from having a file to not.
file.map_or(self.last_full_path.to_path_buf(), |file| file.full_path(cx))
}
}
impl std::fmt::Debug for ContextBuffer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ContextBuffer")
.field("id", &self.id)
.field("buffer", &self.buffer)
.field("version", &self.version)
.field("text", &self.text)
.finish()
}
}
#[derive(Debug, Clone)]
pub struct ContextSymbol {
pub id: ContextSymbolId,
pub buffer: Entity<Buffer>,
pub buffer_version: clock::Global,
/// The range that the symbol encloses, e.g. for function symbol, this will
/// include not only the signature, but also the body
pub enclosing_range: Range<Anchor>,
pub text: SharedString,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ContextSymbolId {
pub path: ProjectPath,
pub name: SharedString,
pub range: Range<Anchor>,
}
#[derive(Debug, Clone)]
pub struct SelectionContext {
pub id: ContextId,
pub range: Range<Anchor>,
pub line_range: Range<Point>,
pub context_buffer: ContextBuffer,
}
#[derive(Debug, Clone)]
pub struct RulesContext {
pub id: ContextId,
pub prompt_id: UserPromptId,
pub title: SharedString,
pub text: SharedString,
}
/// Formats a collection of contexts into a string representation
pub fn format_context_as_string<'a>(
contexts: impl Iterator<Item = &'a AssistantContext>,
cx: &App,
) -> Option<String> {
let mut file_context = Vec::new();
let mut directory_context = Vec::new();
let mut symbol_context = Vec::new();
let mut selection_context = Vec::new();
let mut fetch_context = Vec::new();
let mut thread_context = Vec::new();
let mut rules_context = Vec::new();
for context in contexts {
match context {
AssistantContext::File(context) => file_context.push(context),
AssistantContext::Directory(context) => directory_context.push(context),
AssistantContext::Symbol(context) => symbol_context.push(context),
AssistantContext::Selection(context) => selection_context.push(context),
AssistantContext::FetchedUrl(context) => fetch_context.push(context),
AssistantContext::Thread(context) => thread_context.push(context),
AssistantContext::Rules(context) => rules_context.push(context),
AssistantContext::Image(_) => {}
}
}
if file_context.is_empty()
&& directory_context.is_empty()
&& symbol_context.is_empty()
&& selection_context.is_empty()
&& fetch_context.is_empty()
&& thread_context.is_empty()
&& rules_context.is_empty()
{
return None;
}
let mut result = String::new();
result.push_str("\n<context>\n\
The following items were attached by the user. You don't need to use other tools to read them.\n\n");
if !file_context.is_empty() {
result.push_str("<files>\n");
for context in file_context {
result.push_str(&context.context_buffer.text);
}
result.push_str("</files>\n");
}
if !directory_context.is_empty() {
result.push_str("<directories>\n");
for context in directory_context {
for context_buffer in &context.context_buffers {
result.push_str(&context_buffer.text);
}
}
result.push_str("</directories>\n");
}
if !symbol_context.is_empty() {
result.push_str("<symbols>\n");
for context in symbol_context {
result.push_str(&context.context_symbol.text);
result.push('\n');
}
result.push_str("</symbols>\n");
}
if !selection_context.is_empty() {
result.push_str("<selections>\n");
for context in selection_context {
result.push_str(&context.context_buffer.text);
result.push('\n');
}
result.push_str("</selections>\n");
}
if !fetch_context.is_empty() {
result.push_str("<fetched_urls>\n");
for context in &fetch_context {
result.push_str(&context.url);
result.push('\n');
result.push_str(&context.text);
result.push('\n');
}
result.push_str("</fetched_urls>\n");
}
if !thread_context.is_empty() {
result.push_str("<conversation_threads>\n");
for context in &thread_context {
result.push_str(&context.summary(cx));
result.push('\n');
result.push_str(&context.text);
result.push('\n');
}
result.push_str("</conversation_threads>\n");
}
if !rules_context.is_empty() {
result.push_str(
"<user_rules>\n\
The user has specified the following rules that should be applied:\n\n",
);
for context in &rules_context {
result.push_str(&context.text);
result.push('\n');
}
result.push_str("</user_rules>\n");
}
result.push_str("</context>\n");
Some(result)
}
pub fn attach_context_to_message<'a>(
message: &mut LanguageModelRequestMessage,
contexts: impl Iterator<Item = &'a AssistantContext>,
cx: &App,
) {
if let Some(context_string) = format_context_as_string(contexts, cx) {
message.content.push(context_string.into());
}
}