Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Copy/paste images into editors (Mac only) (#15782)

Browse source 
For future reference: WIP branch of copy/pasting a mixture of images and
text: https://github.com/zed-industries/zed/tree/copy-paste-images -
we'll come back to that one after landing this one.

Release Notes:

- You can now paste images into the Assistant Panel to include them as
context. Currently works only on Mac, and with Anthropic models. Future
support is planned for more models, operating systems, and image
clipboard operations.

---------

Co-authored-by: Antonio <antonio@zed.dev>
Co-authored-by: Mikayla <mikayla@zed.dev>
Co-authored-by: Jason <jason@zed.dev>
Co-authored-by: Kyle <kylek@zed.dev>
This commit is contained in:
Richard Feldman 2024-08-13 13:18:25 -04:00 • committed by GitHub
parent e3b0de5dda
commit b1a581e81b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
58 changed files with 2983 additions and 1708 deletions
Download patch file Download diff file Expand all files Collapse all files

247
crates/assistant/src/context.rs
View file

@ -13,27 +13,31 @@ use editor::Editor;
use fs::{Fs, RemoveOptions};
use futures::{
future::{self, Shared},
stream::FuturesUnordered,
FutureExt, StreamExt,
};
use gpui::{
AppContext, Context as _, EventEmitter, Model, ModelContext, Subscription, Task, UpdateGlobal,
View, WeakView,
AppContext, Context as _, EventEmitter, Image, Model, ModelContext, RenderImage, Subscription,
Task, UpdateGlobal, View, WeakView,
};
use language::{
AnchorRangeExt, Bias, Buffer, BufferSnapshot, LanguageRegistry, OffsetRangeExt, ParseStatus,
Point, ToOffset,
};
use language_model::{
LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, LanguageModelTool,
Role,
LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
LanguageModelTool, Role,
};
use open_ai::Model as OpenAiModel;
use paths::contexts_dir;
use paths::{context_images_dir, contexts_dir};
use project::Project;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use std::{
cmp::{self, Ordering},
collections::hash_map,
fmt::Debug,
iter, mem,
ops::Range,
@ -319,8 +323,23 @@ pub struct MessageMetadata {
timestamp: clock::Lamport,
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug)]
pub struct MessageImage {
image_id: u64,
image: Shared<Task<Option<LanguageModelImage>>>,
}
impl PartialEq for MessageImage {
fn eq(&self, other: &Self) -> bool {
self.image_id == other.image_id
}
}
impl Eq for MessageImage {}
#[derive(Clone, Debug)]
pub struct Message {
pub image_offsets: SmallVec<[(usize, MessageImage); 1]>,
pub offset_range: Range<usize>,
pub index_range: Range<usize>,
pub id: MessageId,
@ -331,13 +350,55 @@ pub struct Message {
impl Message {
fn to_request_message(&self, buffer: &Buffer) -> LanguageModelRequestMessage {
let mut content = Vec::new();
let mut range_start = self.offset_range.start;
for (image_offset, message_image) in self.image_offsets.iter() {
if *image_offset != range_start {
content.push(
buffer
.text_for_range(range_start..*image_offset)
.collect::<String>()
.into(),
)
}
if let Some(image) = message_image.image.clone().now_or_never().flatten() {
content.push(language_model::MessageContent::Image(image));
}
range_start = *image_offset;
}
if range_start != self.offset_range.end {
content.push(
buffer
.text_for_range(range_start..self.offset_range.end)
.collect::<String>()
.into(),
)
}
LanguageModelRequestMessage {
role: self.role,
content: buffer.text_for_range(self.offset_range.clone()).collect(),
content,
}
}
}
#[derive(Clone, Debug)]
pub struct ImageAnchor {
pub anchor: language::Anchor,
pub image_id: u64,
pub render_image: Arc<RenderImage>,
pub image: Shared<Task<Option<LanguageModelImage>>>,
}
impl PartialEq for ImageAnchor {
fn eq(&self, other: &Self) -> bool {
self.image_id == other.image_id
}
}
struct PendingCompletion {
id: usize,
_task: Task<()>,
@ -605,6 +666,8 @@ pub struct Context {
finished_slash_commands: HashSet<SlashCommandId>,
slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,
message_anchors: Vec<MessageAnchor>,
images: HashMap<u64, (Arc<RenderImage>, Shared<Task<Option<LanguageModelImage>>>)>,
image_anchors: Vec<ImageAnchor>,
messages_metadata: HashMap<MessageId, MessageMetadata>,
summary: Option<ContextSummary>,
pending_summary: Task<Option<()>>,
@ -677,6 +740,8 @@ impl Context {
pending_ops: Vec::new(),
operations: Vec::new(),
message_anchors: Default::default(),
image_anchors: Default::default(),
images: Default::default(),
messages_metadata: Default::default(),
pending_slash_commands: Vec::new(),
finished_slash_commands: HashSet::default(),
@ -736,6 +801,11 @@ impl Context {
id: message.id,
start: message.offset_range.start,
metadata: self.messages_metadata[&message.id].clone(),
image_offsets: message
.image_offsets
.iter()
.map(|image_offset| (image_offset.0, image_offset.1.image_id))
.collect(),
})
.collect(),
summary: self
@ -1339,7 +1409,7 @@ impl Context {
request.messages.push(LanguageModelRequestMessage {
role: Role::User,
content: prompt,
content: vec![prompt.into()],
});
// Invoke the model to get its edit suggestions for this workflow step.
@ -1743,13 +1813,15 @@ impl Context {
}
pub fn to_completion_request(&self, cx: &AppContext) -> LanguageModelRequest {
let messages = self
let buffer = self.buffer.read(cx);
let request_messages = self
.messages(cx)
.filter(|message| matches!(message.status, MessageStatus::Done))
.map(|message| message.to_request_message(self.buffer.read(cx)));
.filter(|message| message.status == MessageStatus::Done)
.map(|message| message.to_request_message(&buffer))
.collect();
LanguageModelRequest {
messages: messages.collect(),
messages: request_messages,
stop: vec![],
temperature: 1.0,
}
@ -1847,6 +1919,55 @@ impl Context {
}
}
pub fn insert_image(&mut self, image: Image, cx: &mut ModelContext<Self>) -> Option<()> {
if let hash_map::Entry::Vacant(entry) = self.images.entry(image.id()) {
entry.insert((
image.to_image_data(cx).log_err()?,
LanguageModelImage::from_image(image, cx).shared(),
));
}
Some(())
}
pub fn insert_image_anchor(
&mut self,
image_id: u64,
anchor: language::Anchor,
cx: &mut ModelContext<Self>,
) -> bool {
cx.emit(ContextEvent::MessagesEdited);
let buffer = self.buffer.read(cx);
let insertion_ix = match self
.image_anchors
.binary_search_by(|existing_anchor| anchor.cmp(&existing_anchor.anchor, buffer))
{
Ok(ix) => ix,
Err(ix) => ix,
};
if let Some((render_image, image)) = self.images.get(&image_id) {
self.image_anchors.insert(
insertion_ix,
ImageAnchor {
anchor,
image_id,
image: image.clone(),
render_image: render_image.clone(),
},
);
true
} else {
false
}
}
pub fn images<'a>(&'a self, _cx: &'a AppContext) -> impl 'a + Iterator<Item = ImageAnchor> {
self.image_anchors.iter().cloned()
}
pub fn split_message(
&mut self,
range: Range<usize>,
@ -1865,7 +1986,10 @@ impl Context {
let mut edited_buffer = false;
let mut suffix_start = None;
if range.start > message.offset_range.start && range.end < message.offset_range.end - 1
// TODO: why did this start panicking?
if range.start > message.offset_range.start
&& range.end < message.offset_range.end.saturating_sub(1)
{
if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
suffix_start = Some(range.end + 1);
@ -2007,7 +2131,9 @@ impl Context {
.map(|message| message.to_request_message(self.buffer.read(cx)))
.chain(Some(LanguageModelRequestMessage {
role: Role::User,
content: "Summarize the context into a short title without punctuation.".into(),
content: vec![
"Summarize the context into a short title without punctuation.".into(),
],
}));
let request = LanguageModelRequest {
messages: messages.collect(),
@ -2109,25 +2235,55 @@ impl Context {
pub fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
let buffer = self.buffer.read(cx);
let mut message_anchors = self.message_anchors.iter().enumerate().peekable();
let messages = self.message_anchors.iter().enumerate();
let images = self.image_anchors.iter();
Self::messages_from_iters(buffer, &self.messages_metadata, messages, images)
}
pub fn messages_from_iters<'a>(
buffer: &'a Buffer,
metadata: &'a HashMap<MessageId, MessageMetadata>,
messages: impl Iterator<Item = (usize, &'a MessageAnchor)> + 'a,
images: impl Iterator<Item = &'a ImageAnchor> + 'a,
) -> impl 'a + Iterator<Item = Message> {
let mut messages = messages.peekable();
let mut images = images.peekable();
iter::from_fn(move || {
if let Some((start_ix, message_anchor)) = message_anchors.next() {
let metadata = self.messages_metadata.get(&message_anchor.id)?;
if let Some((start_ix, message_anchor)) = messages.next() {
let metadata = metadata.get(&message_anchor.id)?;
let message_start = message_anchor.start.to_offset(buffer);
let mut message_end = None;
let mut end_ix = start_ix;
while let Some((_, next_message)) = message_anchors.peek() {
while let Some((_, next_message)) = messages.peek() {
if next_message.start.is_valid(buffer) {
message_end = Some(next_message.start);
break;
} else {
end_ix += 1;
message_anchors.next();
messages.next();
}
}
let message_end_anchor = message_end.unwrap_or(language::Anchor::MAX);
let message_end = message_end_anchor.to_offset(buffer);
let mut image_offsets = SmallVec::new();
while let Some(image_anchor) = images.peek() {
if image_anchor.anchor.cmp(&message_end_anchor, buffer).is_lt() {
image_offsets.push((
image_anchor.anchor.to_offset(buffer),
MessageImage {
image_id: image_anchor.image_id,
image: image_anchor.image.clone(),
},
));
images.next();
} else {
break;
}
}
let message_end = message_end
.unwrap_or(language::Anchor::MAX)
.to_offset(buffer);
return Some(Message {
index_range: start_ix..end_ix,
@ -2136,6 +2292,7 @@ impl Context {
anchor: message_anchor.start,
role: metadata.role,
status: metadata.status.clone(),
image_offsets,
});
}
None
@ -2173,6 +2330,9 @@ impl Context {
})?;
if let Some(summary) = summary {
this.read_with(&cx, |this, cx| this.serialize_images(fs.clone(), cx))?
.await;
let context = this.read_with(&cx, |this, cx| this.serialize(cx))?;
let mut discriminant = 1;
let mut new_path;
@ -2212,6 +2372,45 @@ impl Context {
});
}
pub fn serialize_images(&self, fs: Arc<dyn Fs>, cx: &AppContext) -> Task<()> {
let mut images_to_save = self
.images
.iter()
.map(|(id, (_, llm_image))| {
let fs = fs.clone();
let llm_image = llm_image.clone();
let id = *id;
async move {
if let Some(llm_image) = llm_image.await {
let path: PathBuf =
context_images_dir().join(&format!("{}.png.base64", id));
if fs
.metadata(path.as_path())
.await
.log_err()
.flatten()
.is_none()
{
fs.atomic_write(path, llm_image.source.to_string())
.await
.log_err();
}
}
}
})
.collect::<FuturesUnordered<_>>();
cx.background_executor().spawn(async move {
if fs
.create_dir(context_images_dir().as_ref())
.await
.log_err()
.is_some()
{
while let Some(_) = images_to_save.next().await {}
}
})
}
pub(crate) fn custom_summary(&mut self, custom_summary: String, cx: &mut ModelContext<Self>) {
let timestamp = self.next_timestamp();
let summary = self.summary.get_or_insert(ContextSummary::default());
@ -2265,6 +2464,9 @@ pub struct SavedMessage {
pub id: MessageId,
pub start: usize,
pub metadata: MessageMetadata,
#[serde(default)]
// This is defaulted for backwards compatibility with JSON files created before August 2024. We didn't always have this field.
pub image_offsets: Vec<(usize, u64)>,
}
#[derive(Serialize, Deserialize)]
@ -2447,6 +2649,7 @@ impl SavedContextV0_3_0 {
status: metadata.status.clone(),
timestamp,
},
image_offsets: Vec::new(),
})
})
.collect(),