ZIm/crates/ai/src/assistant.rs
Nathan Sobo 1707652643 Always focus a panel when zooming it
This allows us to zoom a panel when clicking a button, even if the
panel isn't currently focused.
2023-06-22 06:55:31 -06:00

2387 lines
87 KiB
Rust

use crate::{
assistant_settings::{AssistantDockPosition, AssistantSettings},
MessageId, MessageMetadata, MessageStatus, OpenAIRequest, OpenAIResponseStreamEvent,
RequestMessage, Role, SavedConversation, SavedConversationMetadata, SavedMessage,
};
use anyhow::{anyhow, Result};
use chrono::{DateTime, Local};
use collections::{HashMap, HashSet};
use editor::{
display_map::{BlockDisposition, BlockId, BlockProperties, BlockStyle, ToDisplayPoint},
scroll::autoscroll::{Autoscroll, AutoscrollStrategy},
Anchor, Editor, ToOffset,
};
use fs::Fs;
use futures::{io::BufReader, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
use gpui::{
actions,
elements::*,
executor::Background,
geometry::vector::{vec2f, Vector2F},
platform::{CursorStyle, MouseButton},
Action, AppContext, AsyncAppContext, ClipboardItem, Entity, ModelContext, ModelHandle,
Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
};
use isahc::{http::StatusCode, Request, RequestExt};
use language::{language_settings::SoftWrap, Buffer, LanguageRegistry, ToOffset as _};
use serde::Deserialize;
use settings::SettingsStore;
use std::{
borrow::Cow,
cell::RefCell,
cmp, env,
fmt::Write,
io, iter,
ops::Range,
path::{Path, PathBuf},
rc::Rc,
sync::Arc,
time::Duration,
};
use theme::{ui::IconStyle, AssistantStyle};
use util::{
channel::ReleaseChannel, paths::CONVERSATIONS_DIR, post_inc, truncate_and_trailoff, ResultExt,
TryFutureExt,
};
use workspace::{
dock::{DockPosition, Panel},
item::Item,
Save, ToggleZoom, Workspace,
};
const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";
actions!(
assistant,
[
NewContext,
Assist,
Split,
CycleMessageRole,
QuoteSelection,
ToggleFocus,
ResetKey,
]
);
pub fn init(cx: &mut AppContext) {
if *util::channel::RELEASE_CHANNEL == ReleaseChannel::Stable {
cx.update_default_global::<collections::CommandPaletteFilter, _, _>(move |filter, _cx| {
filter.filtered_namespaces.insert("assistant");
});
}
settings::register::<AssistantSettings>(cx);
cx.add_action(
|workspace: &mut Workspace, _: &NewContext, cx: &mut ViewContext<Workspace>| {
if let Some(this) = workspace.panel::<AssistantPanel>(cx) {
this.update(cx, |this, cx| {
this.new_conversation(cx);
})
}
workspace.focus_panel::<AssistantPanel>(cx);
},
);
cx.add_action(ConversationEditor::assist);
cx.capture_action(ConversationEditor::cancel_last_assist);
cx.capture_action(ConversationEditor::save);
cx.add_action(ConversationEditor::quote_selection);
cx.capture_action(ConversationEditor::copy);
cx.capture_action(ConversationEditor::split);
cx.capture_action(ConversationEditor::cycle_message_role);
cx.add_action(AssistantPanel::save_api_key);
cx.add_action(AssistantPanel::reset_api_key);
cx.add_action(AssistantPanel::toggle_zoom);
cx.add_action(
|workspace: &mut Workspace, _: &ToggleFocus, cx: &mut ViewContext<Workspace>| {
workspace.toggle_panel_focus::<AssistantPanel>(cx);
},
);
}
#[derive(Debug)]
pub enum AssistantPanelEvent {
ZoomIn,
ZoomOut,
Focus,
Close,
DockPositionChanged,
}
pub struct AssistantPanel {
width: Option<f32>,
height: Option<f32>,
active_editor_index: Option<usize>,
editors: Vec<ViewHandle<ConversationEditor>>,
saved_conversations: Vec<SavedConversationMetadata>,
saved_conversations_list_state: UniformListState,
zoomed: bool,
has_focus: bool,
api_key: Rc<RefCell<Option<String>>>,
api_key_editor: Option<ViewHandle<Editor>>,
has_read_credentials: bool,
languages: Arc<LanguageRegistry>,
fs: Arc<dyn Fs>,
subscriptions: Vec<Subscription>,
_watch_saved_conversations: Task<Result<()>>,
}
impl AssistantPanel {
pub fn load(
workspace: WeakViewHandle<Workspace>,
cx: AsyncAppContext,
) -> Task<Result<ViewHandle<Self>>> {
cx.spawn(|mut cx| async move {
let fs = workspace.read_with(&cx, |workspace, _| workspace.app_state().fs.clone())?;
let saved_conversations = SavedConversationMetadata::list(fs.clone())
.await
.log_err()
.unwrap_or_default();
// TODO: deserialize state.
workspace.update(&mut cx, |workspace, cx| {
cx.add_view::<Self, _>(|cx| {
const CONVERSATION_WATCH_DURATION: Duration = Duration::from_millis(100);
let _watch_saved_conversations = cx.spawn(move |this, mut cx| async move {
let mut events = fs
.watch(&CONVERSATIONS_DIR, CONVERSATION_WATCH_DURATION)
.await;
while events.next().await.is_some() {
let saved_conversations = SavedConversationMetadata::list(fs.clone())
.await
.log_err()
.unwrap_or_default();
this.update(&mut cx, |this, _| {
this.saved_conversations = saved_conversations
})
.ok();
}
anyhow::Ok(())
});
let mut this = Self {
active_editor_index: Default::default(),
editors: Default::default(),
saved_conversations,
saved_conversations_list_state: Default::default(),
zoomed: false,
has_focus: false,
api_key: Rc::new(RefCell::new(None)),
api_key_editor: None,
has_read_credentials: false,
languages: workspace.app_state().languages.clone(),
fs: workspace.app_state().fs.clone(),
width: None,
height: None,
subscriptions: Default::default(),
_watch_saved_conversations,
};
let mut old_dock_position = this.position(cx);
this.subscriptions =
vec![cx.observe_global::<SettingsStore, _>(move |this, cx| {
let new_dock_position = this.position(cx);
if new_dock_position != old_dock_position {
old_dock_position = new_dock_position;
cx.emit(AssistantPanelEvent::DockPositionChanged);
}
})];
this
})
})
})
}
fn new_conversation(&mut self, cx: &mut ViewContext<Self>) -> ViewHandle<ConversationEditor> {
let editor = cx.add_view(|cx| {
ConversationEditor::new(
self.api_key.clone(),
self.languages.clone(),
self.fs.clone(),
cx,
)
});
self.add_conversation(editor.clone(), cx);
editor
}
fn add_conversation(
&mut self,
editor: ViewHandle<ConversationEditor>,
cx: &mut ViewContext<Self>,
) {
self.subscriptions
.push(cx.subscribe(&editor, Self::handle_conversation_editor_event));
let conversation = editor.read(cx).conversation.clone();
self.subscriptions
.push(cx.observe(&conversation, |_, _, cx| cx.notify()));
self.active_editor_index = Some(self.editors.len());
self.editors.push(editor.clone());
if self.has_focus(cx) {
cx.focus(&editor);
}
cx.notify();
}
fn handle_conversation_editor_event(
&mut self,
_: ViewHandle<ConversationEditor>,
event: &ConversationEditorEvent,
cx: &mut ViewContext<Self>,
) {
match event {
ConversationEditorEvent::TabContentChanged => cx.notify(),
}
}
fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
if let Some(api_key) = self
.api_key_editor
.as_ref()
.map(|editor| editor.read(cx).text(cx))
{
if !api_key.is_empty() {
cx.platform()
.write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
.log_err();
*self.api_key.borrow_mut() = Some(api_key);
self.api_key_editor.take();
cx.focus_self();
cx.notify();
}
} else {
cx.propagate_action();
}
}
fn reset_api_key(&mut self, _: &ResetKey, cx: &mut ViewContext<Self>) {
cx.platform().delete_credentials(OPENAI_API_URL).log_err();
self.api_key.take();
self.api_key_editor = Some(build_api_key_editor(cx));
cx.focus_self();
cx.notify();
}
fn toggle_zoom(&mut self, _: &workspace::ToggleZoom, cx: &mut ViewContext<Self>) {
if self.zoomed {
cx.emit(AssistantPanelEvent::ZoomOut)
} else {
cx.emit(AssistantPanelEvent::ZoomIn)
}
}
fn active_editor(&self) -> Option<&ViewHandle<ConversationEditor>> {
self.editors.get(self.active_editor_index?)
}
fn render_hamburger_button(style: &IconStyle) -> impl Element<Self> {
enum ListConversations {}
Svg::for_style(style.icon.clone())
.contained()
.with_style(style.container)
.mouse::<ListConversations>(0)
.with_cursor_style(CursorStyle::PointingHand)
.on_click(MouseButton::Left, |_, this: &mut Self, cx| {
this.active_editor_index = None;
cx.notify();
})
}
fn render_current_model(
&self,
style: &AssistantStyle,
cx: &mut ViewContext<Self>,
) -> Option<impl Element<Self>> {
enum Model {}
let model = self
.active_editor()?
.read(cx)
.conversation
.read(cx)
.model
.clone();
Some(
MouseEventHandler::<Model, _>::new(0, cx, |state, _| {
let style = style.model.style_for(state, false);
Label::new(model, style.text.clone())
.contained()
.with_style(style.container)
})
.with_cursor_style(CursorStyle::PointingHand)
.on_click(MouseButton::Left, |_, this, cx| {
if let Some(editor) = this.active_editor() {
editor.update(cx, |editor, cx| {
editor.cycle_model(cx);
});
}
}),
)
}
fn render_remaining_tokens(
&self,
style: &AssistantStyle,
cx: &mut ViewContext<Self>,
) -> Option<impl Element<Self>> {
self.active_editor().and_then(|editor| {
editor
.read(cx)
.conversation
.read(cx)
.remaining_tokens()
.map(|remaining_tokens| {
let remaining_tokens_style = if remaining_tokens <= 0 {
&style.no_remaining_tokens
} else {
&style.remaining_tokens
};
Label::new(
remaining_tokens.to_string(),
remaining_tokens_style.text.clone(),
)
.contained()
.with_style(remaining_tokens_style.container)
})
})
}
fn render_plus_button(style: &IconStyle) -> impl Element<Self> {
enum AddConversation {}
Svg::for_style(style.icon.clone())
.contained()
.with_style(style.container)
.mouse::<AddConversation>(0)
.with_cursor_style(CursorStyle::PointingHand)
.on_click(MouseButton::Left, |_, this: &mut Self, cx| {
this.new_conversation(cx);
})
}
fn render_zoom_button(
&self,
style: &AssistantStyle,
cx: &mut ViewContext<Self>,
) -> impl Element<Self> {
enum ToggleZoomButton {}
let style = if self.zoomed {
&style.zoom_out_button
} else {
&style.zoom_in_button
};
MouseEventHandler::<ToggleZoomButton, _>::new(0, cx, |_, _| {
Svg::for_style(style.icon.clone())
.contained()
.with_style(style.container)
})
.with_cursor_style(CursorStyle::PointingHand)
.on_click(MouseButton::Left, |_, this, cx| {
this.toggle_zoom(&ToggleZoom, cx);
})
}
fn render_saved_conversation(
&mut self,
index: usize,
cx: &mut ViewContext<Self>,
) -> impl Element<Self> {
let conversation = &self.saved_conversations[index];
let path = conversation.path.clone();
MouseEventHandler::<SavedConversationMetadata, _>::new(index, cx, move |state, cx| {
let style = &theme::current(cx).assistant.saved_conversation;
Flex::row()
.with_child(
Label::new(
conversation.mtime.format("%F %I:%M%p").to_string(),
style.saved_at.text.clone(),
)
.aligned()
.contained()
.with_style(style.saved_at.container),
)
.with_child(
Label::new(conversation.title.clone(), style.title.text.clone())
.aligned()
.contained()
.with_style(style.title.container),
)
.contained()
.with_style(*style.container.style_for(state, false))
})
.with_cursor_style(CursorStyle::PointingHand)
.on_click(MouseButton::Left, move |_, this, cx| {
this.open_conversation(path.clone(), cx)
.detach_and_log_err(cx)
})
}
fn open_conversation(&mut self, path: PathBuf, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
if let Some(ix) = self.editor_index_for_path(&path, cx) {
self.active_editor_index = Some(ix);
cx.notify();
return Task::ready(Ok(()));
}
let fs = self.fs.clone();
let conversation = Conversation::load(
path.clone(),
self.api_key.clone(),
self.languages.clone(),
self.fs.clone(),
cx,
);
cx.spawn(|this, mut cx| async move {
let conversation = conversation.await?;
this.update(&mut cx, |this, cx| {
// If, by the time we've loaded the conversation, the user has already opened
// the same conversation, we don't want to open it again.
if let Some(ix) = this.editor_index_for_path(&path, cx) {
this.active_editor_index = Some(ix);
} else {
let editor = cx
.add_view(|cx| ConversationEditor::from_conversation(conversation, fs, cx));
this.add_conversation(editor, cx);
}
})?;
Ok(())
})
}
fn editor_index_for_path(&self, path: &Path, cx: &AppContext) -> Option<usize> {
self.editors
.iter()
.position(|editor| editor.read(cx).conversation.read(cx).path.as_deref() == Some(path))
}
}
fn build_api_key_editor(cx: &mut ViewContext<AssistantPanel>) -> ViewHandle<Editor> {
cx.add_view(|cx| {
let mut editor = Editor::single_line(
Some(Arc::new(|theme| theme.assistant.api_key_editor.clone())),
cx,
);
editor.set_placeholder_text("sk-000000000000000000000000000000000000000000000000", cx);
editor
})
}
impl Entity for AssistantPanel {
type Event = AssistantPanelEvent;
}
impl View for AssistantPanel {
fn ui_name() -> &'static str {
"AssistantPanel"
}
fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
let theme = &theme::current(cx);
let style = &theme.assistant;
if let Some(api_key_editor) = self.api_key_editor.as_ref() {
Flex::column()
.with_child(
Text::new(
"Paste your OpenAI API key and press Enter to use the assistant",
style.api_key_prompt.text.clone(),
)
.aligned(),
)
.with_child(
ChildView::new(api_key_editor, cx)
.contained()
.with_style(style.api_key_editor.container)
.aligned(),
)
.contained()
.with_style(style.api_key_prompt.container)
.aligned()
.into_any()
} else {
let title = self.active_editor().map(|editor| {
Label::new(editor.read(cx).title(cx), style.title.text.clone())
.contained()
.with_style(style.title.container)
.aligned()
.left()
.flex(1., false)
});
Flex::column()
.with_child(
Flex::row()
.with_child(
Self::render_hamburger_button(&style.hamburger_button).aligned(),
)
.with_children(title)
.with_children(
self.render_current_model(&style, cx)
.map(|current_model| current_model.aligned().flex_float()),
)
.with_children(
self.render_remaining_tokens(&style, cx)
.map(|remaining_tokens| remaining_tokens.aligned().flex_float()),
)
.with_child(
Self::render_plus_button(&style.plus_button)
.aligned()
.flex_float(),
)
.with_child(self.render_zoom_button(&style, cx).aligned().flex_float())
.contained()
.with_style(theme.workspace.tab_bar.container)
.expanded()
.constrained()
.with_height(theme.workspace.tab_bar.height),
)
.with_child(if let Some(editor) = self.active_editor() {
ChildView::new(editor, cx).flex(1., true).into_any()
} else {
UniformList::new(
self.saved_conversations_list_state.clone(),
self.saved_conversations.len(),
cx,
|this, range, items, cx| {
for ix in range {
items.push(this.render_saved_conversation(ix, cx).into_any());
}
},
)
.flex(1., true)
.into_any()
})
.into_any()
}
}
fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
self.has_focus = true;
if cx.is_self_focused() {
if let Some(editor) = self.active_editor() {
cx.focus(editor);
} else if let Some(api_key_editor) = self.api_key_editor.as_ref() {
cx.focus(api_key_editor);
}
}
}
fn focus_out(&mut self, _: gpui::AnyViewHandle, _: &mut ViewContext<Self>) {
self.has_focus = false;
}
}
impl Panel for AssistantPanel {
fn position(&self, cx: &WindowContext) -> DockPosition {
match settings::get::<AssistantSettings>(cx).dock {
AssistantDockPosition::Left => DockPosition::Left,
AssistantDockPosition::Bottom => DockPosition::Bottom,
AssistantDockPosition::Right => DockPosition::Right,
}
}
fn position_is_valid(&self, _: DockPosition) -> bool {
true
}
fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext<Self>) {
settings::update_settings_file::<AssistantSettings>(self.fs.clone(), cx, move |settings| {
let dock = match position {
DockPosition::Left => AssistantDockPosition::Left,
DockPosition::Bottom => AssistantDockPosition::Bottom,
DockPosition::Right => AssistantDockPosition::Right,
};
settings.dock = Some(dock);
});
}
fn size(&self, cx: &WindowContext) -> f32 {
let settings = settings::get::<AssistantSettings>(cx);
match self.position(cx) {
DockPosition::Left | DockPosition::Right => {
self.width.unwrap_or_else(|| settings.default_width)
}
DockPosition::Bottom => self.height.unwrap_or_else(|| settings.default_height),
}
}
fn set_size(&mut self, size: f32, cx: &mut ViewContext<Self>) {
match self.position(cx) {
DockPosition::Left | DockPosition::Right => self.width = Some(size),
DockPosition::Bottom => self.height = Some(size),
}
cx.notify();
}
fn should_zoom_in_on_event(event: &AssistantPanelEvent) -> bool {
matches!(event, AssistantPanelEvent::ZoomIn)
}
fn should_zoom_out_on_event(event: &AssistantPanelEvent) -> bool {
matches!(event, AssistantPanelEvent::ZoomOut)
}
fn is_zoomed(&self, _: &WindowContext) -> bool {
self.zoomed
}
fn set_zoomed(&mut self, zoomed: bool, cx: &mut ViewContext<Self>) {
self.zoomed = zoomed;
cx.notify();
}
fn set_active(&mut self, active: bool, cx: &mut ViewContext<Self>) {
if active {
if self.api_key.borrow().is_none() && !self.has_read_credentials {
self.has_read_credentials = true;
let api_key = if let Ok(api_key) = env::var("OPENAI_API_KEY") {
Some(api_key)
} else if let Some((_, api_key)) = cx
.platform()
.read_credentials(OPENAI_API_URL)
.log_err()
.flatten()
{
String::from_utf8(api_key).log_err()
} else {
None
};
if let Some(api_key) = api_key {
*self.api_key.borrow_mut() = Some(api_key);
} else if self.api_key_editor.is_none() {
self.api_key_editor = Some(build_api_key_editor(cx));
cx.notify();
}
}
if self.editors.is_empty() {
self.new_conversation(cx);
}
}
}
fn icon_path(&self) -> &'static str {
"icons/robot_14.svg"
}
fn icon_tooltip(&self) -> (String, Option<Box<dyn Action>>) {
("Assistant Panel".into(), Some(Box::new(ToggleFocus)))
}
fn should_change_position_on_event(event: &Self::Event) -> bool {
matches!(event, AssistantPanelEvent::DockPositionChanged)
}
fn should_activate_on_event(_: &Self::Event) -> bool {
false
}
fn should_close_on_event(event: &AssistantPanelEvent) -> bool {
matches!(event, AssistantPanelEvent::Close)
}
fn has_focus(&self, _: &WindowContext) -> bool {
self.has_focus
}
fn is_focus_event(event: &Self::Event) -> bool {
matches!(event, AssistantPanelEvent::Focus)
}
}
enum ConversationEvent {
MessagesEdited,
SummaryChanged,
StreamedCompletion,
}
#[derive(Default)]
struct Summary {
text: String,
done: bool,
}
struct Conversation {
buffer: ModelHandle<Buffer>,
message_anchors: Vec<MessageAnchor>,
messages_metadata: HashMap<MessageId, MessageMetadata>,
next_message_id: MessageId,
summary: Option<Summary>,
pending_summary: Task<Option<()>>,
completion_count: usize,
pending_completions: Vec<PendingCompletion>,
model: String,
token_count: Option<usize>,
max_token_count: usize,
pending_token_count: Task<Option<()>>,
api_key: Rc<RefCell<Option<String>>>,
pending_save: Task<Result<()>>,
path: Option<PathBuf>,
_subscriptions: Vec<Subscription>,
}
impl Entity for Conversation {
type Event = ConversationEvent;
}
impl Conversation {
fn new(
api_key: Rc<RefCell<Option<String>>>,
language_registry: Arc<LanguageRegistry>,
cx: &mut ModelContext<Self>,
) -> Self {
let model = "gpt-3.5-turbo-0613";
let markdown = language_registry.language_for_name("Markdown");
let buffer = cx.add_model(|cx| {
let mut buffer = Buffer::new(0, "", cx);
buffer.set_language_registry(language_registry);
cx.spawn_weak(|buffer, mut cx| async move {
let markdown = markdown.await?;
let buffer = buffer
.upgrade(&cx)
.ok_or_else(|| anyhow!("buffer was dropped"))?;
buffer.update(&mut cx, |buffer, cx| {
buffer.set_language(Some(markdown), cx)
});
anyhow::Ok(())
})
.detach_and_log_err(cx);
buffer
});
let mut this = Self {
message_anchors: Default::default(),
messages_metadata: Default::default(),
next_message_id: Default::default(),
summary: None,
pending_summary: Task::ready(None),
completion_count: Default::default(),
pending_completions: Default::default(),
token_count: None,
max_token_count: tiktoken_rs::model::get_context_size(model),
pending_token_count: Task::ready(None),
model: model.into(),
_subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
pending_save: Task::ready(Ok(())),
path: None,
api_key,
buffer,
};
let message = MessageAnchor {
id: MessageId(post_inc(&mut this.next_message_id.0)),
start: language::Anchor::MIN,
};
this.message_anchors.push(message.clone());
this.messages_metadata.insert(
message.id,
MessageMetadata {
role: Role::User,
sent_at: Local::now(),
status: MessageStatus::Done,
},
);
this.count_remaining_tokens(cx);
this
}
fn load(
path: PathBuf,
api_key: Rc<RefCell<Option<String>>>,
language_registry: Arc<LanguageRegistry>,
fs: Arc<dyn Fs>,
cx: &mut AppContext,
) -> Task<Result<ModelHandle<Self>>> {
cx.spawn(|mut cx| async move {
let saved_conversation = fs.load(&path).await?;
let saved_conversation: SavedConversation = serde_json::from_str(&saved_conversation)?;
let model = saved_conversation.model;
let markdown = language_registry.language_for_name("Markdown");
let mut message_anchors = Vec::new();
let mut next_message_id = MessageId(0);
let buffer = cx.add_model(|cx| {
let mut buffer = Buffer::new(0, saved_conversation.text, cx);
for message in saved_conversation.messages {
message_anchors.push(MessageAnchor {
id: message.id,
start: buffer.anchor_before(message.start),
});
next_message_id = cmp::max(next_message_id, MessageId(message.id.0 + 1));
}
buffer.set_language_registry(language_registry);
cx.spawn_weak(|buffer, mut cx| async move {
let markdown = markdown.await?;
let buffer = buffer
.upgrade(&cx)
.ok_or_else(|| anyhow!("buffer was dropped"))?;
buffer.update(&mut cx, |buffer, cx| {
buffer.set_language(Some(markdown), cx)
});
anyhow::Ok(())
})
.detach_and_log_err(cx);
buffer
});
let conversation = cx.add_model(|cx| {
let mut this = Self {
message_anchors,
messages_metadata: saved_conversation.message_metadata,
next_message_id,
summary: Some(Summary {
text: saved_conversation.summary,
done: true,
}),
pending_summary: Task::ready(None),
completion_count: Default::default(),
pending_completions: Default::default(),
token_count: None,
max_token_count: tiktoken_rs::model::get_context_size(&model),
pending_token_count: Task::ready(None),
model,
_subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
pending_save: Task::ready(Ok(())),
path: Some(path),
api_key,
buffer,
};
this.count_remaining_tokens(cx);
this
});
Ok(conversation)
})
}
fn handle_buffer_event(
&mut self,
_: ModelHandle<Buffer>,
event: &language::Event,
cx: &mut ModelContext<Self>,
) {
match event {
language::Event::Edited => {
self.count_remaining_tokens(cx);
cx.emit(ConversationEvent::MessagesEdited);
}
_ => {}
}
}
fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
let messages = self
.messages(cx)
.into_iter()
.filter_map(|message| {
Some(tiktoken_rs::ChatCompletionRequestMessage {
role: match message.role {
Role::User => "user".into(),
Role::Assistant => "assistant".into(),
Role::System => "system".into(),
},
content: self.buffer.read(cx).text_for_range(message.range).collect(),
name: None,
})
})
.collect::<Vec<_>>();
let model = self.model.clone();
self.pending_token_count = cx.spawn_weak(|this, mut cx| {
async move {
cx.background().timer(Duration::from_millis(200)).await;
let token_count = cx
.background()
.spawn(async move { tiktoken_rs::num_tokens_from_messages(&model, &messages) })
.await?;
this.upgrade(&cx)
.ok_or_else(|| anyhow!("conversation was dropped"))?
.update(&mut cx, |this, cx| {
this.max_token_count = tiktoken_rs::model::get_context_size(&this.model);
this.token_count = Some(token_count);
cx.notify()
});
anyhow::Ok(())
}
.log_err()
});
}
fn remaining_tokens(&self) -> Option<isize> {
Some(self.max_token_count as isize - self.token_count? as isize)
}
fn set_model(&mut self, model: String, cx: &mut ModelContext<Self>) {
self.model = model;
self.count_remaining_tokens(cx);
cx.notify();
}
fn assist(
&mut self,
selected_messages: HashSet<MessageId>,
cx: &mut ModelContext<Self>,
) -> Vec<MessageAnchor> {
let mut user_messages = Vec::new();
let mut tasks = Vec::new();
for selected_message_id in selected_messages {
let selected_message_role =
if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
metadata.role
} else {
continue;
};
if selected_message_role == Role::Assistant {
if let Some(user_message) = self.insert_message_after(
selected_message_id,
Role::User,
MessageStatus::Done,
cx,
) {
user_messages.push(user_message);
} else {
continue;
}
} else {
let request = OpenAIRequest {
model: self.model.clone(),
messages: self
.messages(cx)
.filter(|message| matches!(message.status, MessageStatus::Done))
.flat_map(|message| {
let mut system_message = None;
if message.id == selected_message_id {
system_message = Some(RequestMessage {
role: Role::System,
content: concat!(
"Treat the following messages as additional knowledge you have learned about, ",
"but act as if they were not part of this conversation. That is, treat them ",
"as if the user didn't see them and couldn't possibly inquire about them."
).into()
});
}
Some(message.to_open_ai_message(self.buffer.read(cx))).into_iter().chain(system_message)
})
.chain(Some(RequestMessage {
role: Role::System,
content: format!(
"Direct your reply to message with id {}. Do not include a [Message X] header.",
selected_message_id.0
),
}))
.collect(),
stream: true,
};
let Some(api_key) = self.api_key.borrow().clone() else { continue };
let stream = stream_completion(api_key, cx.background().clone(), request);
let assistant_message = self
.insert_message_after(
selected_message_id,
Role::Assistant,
MessageStatus::Pending,
cx,
)
.unwrap();
tasks.push(cx.spawn_weak({
|this, mut cx| async move {
let assistant_message_id = assistant_message.id;
let stream_completion = async {
let mut messages = stream.await?;
while let Some(message) = messages.next().await {
let mut message = message?;
if let Some(choice) = message.choices.pop() {
this.upgrade(&cx)
.ok_or_else(|| anyhow!("conversation was dropped"))?
.update(&mut cx, |this, cx| {
let text: Arc<str> = choice.delta.content?.into();
let message_ix = this.message_anchors.iter().position(
|message| message.id == assistant_message_id,
)?;
this.buffer.update(cx, |buffer, cx| {
let offset = this.message_anchors[message_ix + 1..]
.iter()
.find(|message| message.start.is_valid(buffer))
.map_or(buffer.len(), |message| {
message
.start
.to_offset(buffer)
.saturating_sub(1)
});
buffer.edit([(offset..offset, text)], None, cx);
});
cx.emit(ConversationEvent::StreamedCompletion);
Some(())
});
}
smol::future::yield_now().await;
}
this.upgrade(&cx)
.ok_or_else(|| anyhow!("conversation was dropped"))?
.update(&mut cx, |this, cx| {
this.pending_completions.retain(|completion| {
completion.id != this.completion_count
});
this.summarize(cx);
});
anyhow::Ok(())
};
let result = stream_completion.await;
if let Some(this) = this.upgrade(&cx) {
this.update(&mut cx, |this, cx| {
if let Some(metadata) =
this.messages_metadata.get_mut(&assistant_message.id)
{
match result {
Ok(_) => {
metadata.status = MessageStatus::Done;
}
Err(error) => {
metadata.status = MessageStatus::Error(
error.to_string().trim().into(),
);
}
}
cx.notify();
}
});
}
}
}));
}
}
if !tasks.is_empty() {
self.pending_completions.push(PendingCompletion {
id: post_inc(&mut self.completion_count),
_tasks: tasks,
});
}
user_messages
}
fn cancel_last_assist(&mut self) -> bool {
self.pending_completions.pop().is_some()
}
fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
for id in ids {
if let Some(metadata) = self.messages_metadata.get_mut(&id) {
metadata.role.cycle();
cx.emit(ConversationEvent::MessagesEdited);
cx.notify();
}
}
}
fn insert_message_after(
&mut self,
message_id: MessageId,
role: Role,
status: MessageStatus,
cx: &mut ModelContext<Self>,
) -> Option<MessageAnchor> {
if let Some(prev_message_ix) = self
.message_anchors
.iter()
.position(|message| message.id == message_id)
{
let start = self.buffer.update(cx, |buffer, cx| {
let offset = self.message_anchors[prev_message_ix + 1..]
.iter()
.find(|message| message.start.is_valid(buffer))
.map_or(buffer.len(), |message| message.start.to_offset(buffer) - 1);
buffer.edit([(offset..offset, "\n")], None, cx);
buffer.anchor_before(offset + 1)
});
let message = MessageAnchor {
id: MessageId(post_inc(&mut self.next_message_id.0)),
start,
};
self.message_anchors
.insert(prev_message_ix + 1, message.clone());
self.messages_metadata.insert(
message.id,
MessageMetadata {
role,
sent_at: Local::now(),
status,
},
);
cx.emit(ConversationEvent::MessagesEdited);
Some(message)
} else {
None
}
}
fn split_message(
&mut self,
range: Range<usize>,
cx: &mut ModelContext<Self>,
) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
let start_message = self.message_for_offset(range.start, cx);
let end_message = self.message_for_offset(range.end, cx);
if let Some((start_message, end_message)) = start_message.zip(end_message) {
// Prevent splitting when range spans multiple messages.
if start_message.index != end_message.index {
return (None, None);
}
let message = start_message;
let role = message.role;
let mut edited_buffer = false;
let mut suffix_start = None;
if range.start > message.range.start && range.end < message.range.end - 1 {
if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
suffix_start = Some(range.end + 1);
} else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
suffix_start = Some(range.end);
}
}
let suffix = if let Some(suffix_start) = suffix_start {
MessageAnchor {
id: MessageId(post_inc(&mut self.next_message_id.0)),
start: self.buffer.read(cx).anchor_before(suffix_start),
}
} else {
self.buffer.update(cx, |buffer, cx| {
buffer.edit([(range.end..range.end, "\n")], None, cx);
});
edited_buffer = true;
MessageAnchor {
id: MessageId(post_inc(&mut self.next_message_id.0)),
start: self.buffer.read(cx).anchor_before(range.end + 1),
}
};
self.message_anchors
.insert(message.index + 1, suffix.clone());
self.messages_metadata.insert(
suffix.id,
MessageMetadata {
role,
sent_at: Local::now(),
status: MessageStatus::Done,
},
);
let new_messages = if range.start == range.end || range.start == message.range.start {
(None, Some(suffix))
} else {
let mut prefix_end = None;
if range.start > message.range.start && range.end < message.range.end - 1 {
if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
prefix_end = Some(range.start + 1);
} else if self.buffer.read(cx).reversed_chars_at(range.start).next()
== Some('\n')
{
prefix_end = Some(range.start);
}
}
let selection = if let Some(prefix_end) = prefix_end {
cx.emit(ConversationEvent::MessagesEdited);
MessageAnchor {
id: MessageId(post_inc(&mut self.next_message_id.0)),
start: self.buffer.read(cx).anchor_before(prefix_end),
}
} else {
self.buffer.update(cx, |buffer, cx| {
buffer.edit([(range.start..range.start, "\n")], None, cx)
});
edited_buffer = true;
MessageAnchor {
id: MessageId(post_inc(&mut self.next_message_id.0)),
start: self.buffer.read(cx).anchor_before(range.end + 1),
}
};
self.message_anchors
.insert(message.index + 1, selection.clone());
self.messages_metadata.insert(
selection.id,
MessageMetadata {
role,
sent_at: Local::now(),
status: MessageStatus::Done,
},
);
(Some(selection), Some(suffix))
};
if !edited_buffer {
cx.emit(ConversationEvent::MessagesEdited);
}
new_messages
} else {
(None, None)
}
}
fn summarize(&mut self, cx: &mut ModelContext<Self>) {
if self.message_anchors.len() >= 2 && self.summary.is_none() {
let api_key = self.api_key.borrow().clone();
if let Some(api_key) = api_key {
let messages = self
.messages(cx)
.take(2)
.map(|message| message.to_open_ai_message(self.buffer.read(cx)))
.chain(Some(RequestMessage {
role: Role::User,
content:
"Summarize the conversation into a short title without punctuation"
.into(),
}));
let request = OpenAIRequest {
model: self.model.clone(),
messages: messages.collect(),
stream: true,
};
let stream = stream_completion(api_key, cx.background().clone(), request);
self.pending_summary = cx.spawn(|this, mut cx| {
async move {
let mut messages = stream.await?;
while let Some(message) = messages.next().await {
let mut message = message?;
if let Some(choice) = message.choices.pop() {
let text = choice.delta.content.unwrap_or_default();
this.update(&mut cx, |this, cx| {
this.summary
.get_or_insert(Default::default())
.text
.push_str(&text);
cx.emit(ConversationEvent::SummaryChanged);
});
}
}
this.update(&mut cx, |this, cx| {
if let Some(summary) = this.summary.as_mut() {
summary.done = true;
cx.emit(ConversationEvent::SummaryChanged);
}
});
anyhow::Ok(())
}
.log_err()
});
}
}
}
fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
self.messages_for_offsets([offset], cx).pop()
}
fn messages_for_offsets(
&self,
offsets: impl IntoIterator<Item = usize>,
cx: &AppContext,
) -> Vec<Message> {
let mut result = Vec::new();
let buffer_len = self.buffer.read(cx).len();
let mut messages = self.messages(cx).peekable();
let mut offsets = offsets.into_iter().peekable();
while let Some(offset) = offsets.next() {
// Skip messages that start after the offset.
while messages.peek().map_or(false, |message| {
message.range.end < offset || (message.range.end == offset && offset < buffer_len)
}) {
messages.next();
}
let Some(message) = messages.peek() else { continue };
// Skip offsets that are in the same message.
while offsets.peek().map_or(false, |offset| {
message.range.contains(offset) || message.range.end == buffer_len
}) {
offsets.next();
}
result.push(message.clone());
}
result
}
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();
iter::from_fn(move || {
while let Some((ix, message_anchor)) = message_anchors.next() {
let metadata = self.messages_metadata.get(&message_anchor.id)?;
let message_start = message_anchor.start.to_offset(buffer);
let mut message_end = None;
while let Some((_, next_message)) = message_anchors.peek() {
if next_message.start.is_valid(buffer) {
message_end = Some(next_message.start);
break;
} else {
message_anchors.next();
}
}
let message_end = message_end
.unwrap_or(language::Anchor::MAX)
.to_offset(buffer);
return Some(Message {
index: ix,
range: message_start..message_end,
id: message_anchor.id,
anchor: message_anchor.start,
role: metadata.role,
sent_at: metadata.sent_at,
status: metadata.status.clone(),
});
}
None
})
}
fn save(
&mut self,
debounce: Option<Duration>,
fs: Arc<dyn Fs>,
cx: &mut ModelContext<Conversation>,
) {
self.pending_save = cx.spawn(|this, mut cx| async move {
if let Some(debounce) = debounce {
cx.background().timer(debounce).await;
}
let (old_path, summary) = this.read_with(&cx, |this, _| {
let path = this.path.clone();
let summary = if let Some(summary) = this.summary.as_ref() {
if summary.done {
Some(summary.text.clone())
} else {
None
}
} else {
None
};
(path, summary)
});
if let Some(summary) = summary {
let conversation = this.read_with(&cx, |this, cx| SavedConversation {
zed: "conversation".into(),
version: SavedConversation::VERSION.into(),
text: this.buffer.read(cx).text(),
message_metadata: this.messages_metadata.clone(),
messages: this
.message_anchors
.iter()
.map(|message| SavedMessage {
id: message.id,
start: message.start.to_offset(this.buffer.read(cx)),
})
.collect(),
summary: summary.clone(),
model: this.model.clone(),
});
let path = if let Some(old_path) = old_path {
old_path
} else {
let mut discriminant = 1;
let mut new_path;
loop {
new_path = CONVERSATIONS_DIR.join(&format!(
"{} - {}.zed.json",
summary.trim(),
discriminant
));
if fs.is_file(&new_path).await {
discriminant += 1;
} else {
break;
}
}
new_path
};
fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?;
fs.atomic_write(path.clone(), serde_json::to_string(&conversation).unwrap())
.await?;
this.update(&mut cx, |this, _| this.path = Some(path));
}
Ok(())
});
}
}
struct PendingCompletion {
id: usize,
_tasks: Vec<Task<()>>,
}
enum ConversationEditorEvent {
TabContentChanged,
}
#[derive(Copy, Clone, Debug, PartialEq)]
struct ScrollPosition {
offset_before_cursor: Vector2F,
cursor: Anchor,
}
struct ConversationEditor {
conversation: ModelHandle<Conversation>,
fs: Arc<dyn Fs>,
editor: ViewHandle<Editor>,
blocks: HashSet<BlockId>,
scroll_position: Option<ScrollPosition>,
_subscriptions: Vec<Subscription>,
}
impl ConversationEditor {
fn new(
api_key: Rc<RefCell<Option<String>>>,
language_registry: Arc<LanguageRegistry>,
fs: Arc<dyn Fs>,
cx: &mut ViewContext<Self>,
) -> Self {
let conversation = cx.add_model(|cx| Conversation::new(api_key, language_registry, cx));
Self::from_conversation(conversation, fs, cx)
}
fn from_conversation(
conversation: ModelHandle<Conversation>,
fs: Arc<dyn Fs>,
cx: &mut ViewContext<Self>,
) -> Self {
let editor = cx.add_view(|cx| {
let mut editor = Editor::for_buffer(conversation.read(cx).buffer.clone(), None, cx);
editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
editor.set_show_gutter(false, cx);
editor
});
let _subscriptions = vec![
cx.observe(&conversation, |_, _, cx| cx.notify()),
cx.subscribe(&conversation, Self::handle_conversation_event),
cx.subscribe(&editor, Self::handle_editor_event),
];
let mut this = Self {
conversation,
editor,
blocks: Default::default(),
scroll_position: None,
fs,
_subscriptions,
};
this.update_message_headers(cx);
this
}
fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
let cursors = self.cursors(cx);
let user_messages = self.conversation.update(cx, |conversation, cx| {
let selected_messages = conversation
.messages_for_offsets(cursors, cx)
.into_iter()
.map(|message| message.id)
.collect();
conversation.assist(selected_messages, cx)
});
let new_selections = user_messages
.iter()
.map(|message| {
let cursor = message
.start
.to_offset(self.conversation.read(cx).buffer.read(cx));
cursor..cursor
})
.collect::<Vec<_>>();
if !new_selections.is_empty() {
self.editor.update(cx, |editor, cx| {
editor.change_selections(
Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
cx,
|selections| selections.select_ranges(new_selections),
);
});
}
}
fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
if !self
.conversation
.update(cx, |conversation, _| conversation.cancel_last_assist())
{
cx.propagate_action();
}
}
fn cycle_message_role(&mut self, _: &CycleMessageRole, cx: &mut ViewContext<Self>) {
let cursors = self.cursors(cx);
self.conversation.update(cx, |conversation, cx| {
let messages = conversation
.messages_for_offsets(cursors, cx)
.into_iter()
.map(|message| message.id)
.collect();
conversation.cycle_message_roles(messages, cx)
});
}
fn cursors(&self, cx: &AppContext) -> Vec<usize> {
let selections = self.editor.read(cx).selections.all::<usize>(cx);
selections
.into_iter()
.map(|selection| selection.head())
.collect()
}
fn handle_conversation_event(
&mut self,
_: ModelHandle<Conversation>,
event: &ConversationEvent,
cx: &mut ViewContext<Self>,
) {
match event {
ConversationEvent::MessagesEdited => {
self.update_message_headers(cx);
self.conversation.update(cx, |conversation, cx| {
conversation.save(Some(Duration::from_millis(500)), self.fs.clone(), cx);
});
}
ConversationEvent::SummaryChanged => {
cx.emit(ConversationEditorEvent::TabContentChanged);
self.conversation.update(cx, |conversation, cx| {
conversation.save(None, self.fs.clone(), cx);
});
}
ConversationEvent::StreamedCompletion => {
self.editor.update(cx, |editor, cx| {
if let Some(scroll_position) = self.scroll_position {
let snapshot = editor.snapshot(cx);
let cursor_point = scroll_position.cursor.to_display_point(&snapshot);
let scroll_top =
cursor_point.row() as f32 - scroll_position.offset_before_cursor.y();
editor.set_scroll_position(
vec2f(scroll_position.offset_before_cursor.x(), scroll_top),
cx,
);
}
});
}
}
}
fn handle_editor_event(
&mut self,
_: ViewHandle<Editor>,
event: &editor::Event,
cx: &mut ViewContext<Self>,
) {
match event {
editor::Event::ScrollPositionChanged { autoscroll, .. } => {
let cursor_scroll_position = self.cursor_scroll_position(cx);
if *autoscroll {
self.scroll_position = cursor_scroll_position;
} else if self.scroll_position != cursor_scroll_position {
self.scroll_position = None;
}
}
editor::Event::SelectionsChanged { .. } => {
self.scroll_position = self.cursor_scroll_position(cx);
}
_ => {}
}
}
fn cursor_scroll_position(&self, cx: &mut ViewContext<Self>) -> Option<ScrollPosition> {
self.editor.update(cx, |editor, cx| {
let snapshot = editor.snapshot(cx);
let cursor = editor.selections.newest_anchor().head();
let cursor_row = cursor.to_display_point(&snapshot.display_snapshot).row() as f32;
let scroll_position = editor
.scroll_manager
.anchor()
.scroll_position(&snapshot.display_snapshot);
let scroll_bottom = scroll_position.y() + editor.visible_line_count().unwrap_or(0.);
if (scroll_position.y()..scroll_bottom).contains(&cursor_row) {
Some(ScrollPosition {
cursor,
offset_before_cursor: vec2f(
scroll_position.x(),
cursor_row - scroll_position.y(),
),
})
} else {
None
}
})
}
fn update_message_headers(&mut self, cx: &mut ViewContext<Self>) {
self.editor.update(cx, |editor, cx| {
let buffer = editor.buffer().read(cx).snapshot(cx);
let excerpt_id = *buffer.as_singleton().unwrap().0;
let old_blocks = std::mem::take(&mut self.blocks);
let new_blocks = self
.conversation
.read(cx)
.messages(cx)
.map(|message| BlockProperties {
position: buffer.anchor_in_excerpt(excerpt_id, message.anchor),
height: 2,
style: BlockStyle::Sticky,
render: Arc::new({
let conversation = self.conversation.clone();
// let metadata = message.metadata.clone();
// let message = message.clone();
move |cx| {
enum Sender {}
enum ErrorTooltip {}
let theme = theme::current(cx);
let style = &theme.assistant;
let message_id = message.id;
let sender = MouseEventHandler::<Sender, _>::new(
message_id.0,
cx,
|state, _| match message.role {
Role::User => {
let style = style.user_sender.style_for(state, false);
Label::new("You", style.text.clone())
.contained()
.with_style(style.container)
}
Role::Assistant => {
let style = style.assistant_sender.style_for(state, false);
Label::new("Assistant", style.text.clone())
.contained()
.with_style(style.container)
}
Role::System => {
let style = style.system_sender.style_for(state, false);
Label::new("System", style.text.clone())
.contained()
.with_style(style.container)
}
},
)
.with_cursor_style(CursorStyle::PointingHand)
.on_down(MouseButton::Left, {
let conversation = conversation.clone();
move |_, _, cx| {
conversation.update(cx, |conversation, cx| {
conversation.cycle_message_roles(
HashSet::from_iter(Some(message_id)),
cx,
)
})
}
});
Flex::row()
.with_child(sender.aligned())
.with_child(
Label::new(
message.sent_at.format("%I:%M%P").to_string(),
style.sent_at.text.clone(),
)
.contained()
.with_style(style.sent_at.container)
.aligned(),
)
.with_children(
if let MessageStatus::Error(error) = &message.status {
Some(
Svg::new("icons/circle_x_mark_12.svg")
.with_color(style.error_icon.color)
.constrained()
.with_width(style.error_icon.width)
.contained()
.with_style(style.error_icon.container)
.with_tooltip::<ErrorTooltip>(
message_id.0,
error.to_string(),
None,
theme.tooltip.clone(),
cx,
)
.aligned(),
)
} else {
None
},
)
.aligned()
.left()
.contained()
.with_style(style.message_header)
.into_any()
}
}),
disposition: BlockDisposition::Above,
})
.collect::<Vec<_>>();
editor.remove_blocks(old_blocks, None, cx);
let ids = editor.insert_blocks(new_blocks, None, cx);
self.blocks = HashSet::from_iter(ids);
});
}
fn quote_selection(
workspace: &mut Workspace,
_: &QuoteSelection,
cx: &mut ViewContext<Workspace>,
) {
let Some(panel) = workspace.panel::<AssistantPanel>(cx) else {
return;
};
let Some(editor) = workspace.active_item(cx).and_then(|item| item.downcast::<Editor>()) else {
return;
};
let text = editor.read_with(cx, |editor, cx| {
let range = editor.selections.newest::<usize>(cx).range();
let buffer = editor.buffer().read(cx).snapshot(cx);
let start_language = buffer.language_at(range.start);
let end_language = buffer.language_at(range.end);
let language_name = if start_language == end_language {
start_language.map(|language| language.name())
} else {
None
};
let language_name = language_name.as_deref().unwrap_or("").to_lowercase();
let selected_text = buffer.text_for_range(range).collect::<String>();
if selected_text.is_empty() {
None
} else {
Some(if language_name == "markdown" {
selected_text
.lines()
.map(|line| format!("> {}", line))
.collect::<Vec<_>>()
.join("\n")
} else {
format!("```{language_name}\n{selected_text}\n```")
})
}
});
// Activate the panel
if !panel.read(cx).has_focus(cx) {
workspace.toggle_panel_focus::<AssistantPanel>(cx);
}
if let Some(text) = text {
panel.update(cx, |panel, cx| {
let conversation = panel
.active_editor()
.cloned()
.unwrap_or_else(|| panel.new_conversation(cx));
conversation.update(cx, |conversation, cx| {
conversation
.editor
.update(cx, |editor, cx| editor.insert(&text, cx))
});
});
}
}
fn copy(&mut self, _: &editor::Copy, cx: &mut ViewContext<Self>) {
let editor = self.editor.read(cx);
let conversation = self.conversation.read(cx);
if editor.selections.count() == 1 {
let selection = editor.selections.newest::<usize>(cx);
let mut copied_text = String::new();
let mut spanned_messages = 0;
for message in conversation.messages(cx) {
if message.range.start >= selection.range().end {
break;
} else if message.range.end >= selection.range().start {
let range = cmp::max(message.range.start, selection.range().start)
..cmp::min(message.range.end, selection.range().end);
if !range.is_empty() {
spanned_messages += 1;
write!(&mut copied_text, "## {}\n\n", message.role).unwrap();
for chunk in conversation.buffer.read(cx).text_for_range(range) {
copied_text.push_str(&chunk);
}
copied_text.push('\n');
}
}
}
if spanned_messages > 1 {
cx.platform()
.write_to_clipboard(ClipboardItem::new(copied_text));
return;
}
}
cx.propagate_action();
}
fn split(&mut self, _: &Split, cx: &mut ViewContext<Self>) {
self.conversation.update(cx, |conversation, cx| {
let selections = self.editor.read(cx).selections.disjoint_anchors();
for selection in selections.into_iter() {
let buffer = self.editor.read(cx).buffer().read(cx).snapshot(cx);
let range = selection
.map(|endpoint| endpoint.to_offset(&buffer))
.range();
conversation.split_message(range, cx);
}
});
}
fn save(&mut self, _: &Save, cx: &mut ViewContext<Self>) {
self.conversation.update(cx, |conversation, cx| {
conversation.save(None, self.fs.clone(), cx)
});
}
fn cycle_model(&mut self, cx: &mut ViewContext<Self>) {
self.conversation.update(cx, |conversation, cx| {
let new_model = match conversation.model.as_str() {
"gpt-4-0613" => "gpt-3.5-turbo-0613",
_ => "gpt-4-0613",
};
conversation.set_model(new_model.into(), cx);
});
}
fn title(&self, cx: &AppContext) -> String {
self.conversation
.read(cx)
.summary
.as_ref()
.map(|summary| summary.text.clone())
.unwrap_or_else(|| "New Conversation".into())
}
}
impl Entity for ConversationEditor {
type Event = ConversationEditorEvent;
}
impl View for ConversationEditor {
fn ui_name() -> &'static str {
"ConversationEditor"
}
fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
let theme = &theme::current(cx).assistant;
ChildView::new(&self.editor, cx)
.contained()
.with_style(theme.container)
.into_any()
}
fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
if cx.is_self_focused() {
cx.focus(&self.editor);
}
}
}
impl Item for ConversationEditor {
fn tab_content<V: View>(
&self,
_: Option<usize>,
style: &theme::Tab,
cx: &gpui::AppContext,
) -> AnyElement<V> {
let title = truncate_and_trailoff(&self.title(cx), editor::MAX_TAB_TITLE_LEN);
Label::new(title, style.label.clone()).into_any()
}
fn tab_tooltip_text(&self, cx: &AppContext) -> Option<Cow<str>> {
Some(self.title(cx).into())
}
fn as_searchable(
&self,
_: &ViewHandle<Self>,
) -> Option<Box<dyn workspace::searchable::SearchableItemHandle>> {
Some(Box::new(self.editor.clone()))
}
}
#[derive(Clone, Debug)]
struct MessageAnchor {
id: MessageId,
start: language::Anchor,
}
#[derive(Clone, Debug)]
pub struct Message {
range: Range<usize>,
index: usize,
id: MessageId,
anchor: language::Anchor,
role: Role,
sent_at: DateTime<Local>,
status: MessageStatus,
}
impl Message {
fn to_open_ai_message(&self, buffer: &Buffer) -> RequestMessage {
let mut content = format!("[Message {}]\n", self.id.0).to_string();
content.extend(buffer.text_for_range(self.range.clone()));
RequestMessage {
role: self.role,
content: content.trim_end().into(),
}
}
}
async fn stream_completion(
api_key: String,
executor: Arc<Background>,
mut request: OpenAIRequest,
) -> Result<impl Stream<Item = Result<OpenAIResponseStreamEvent>>> {
request.stream = true;
let (tx, rx) = futures::channel::mpsc::unbounded::<Result<OpenAIResponseStreamEvent>>();
let json_data = serde_json::to_string(&request)?;
let mut response = Request::post(format!("{OPENAI_API_URL}/chat/completions"))
.header("Content-Type", "application/json")
.header("Authorization", format!("Bearer {}", api_key))
.body(json_data)?
.send_async()
.await?;
let status = response.status();
if status == StatusCode::OK {
executor
.spawn(async move {
let mut lines = BufReader::new(response.body_mut()).lines();
fn parse_line(
line: Result<String, io::Error>,
) -> Result<Option<OpenAIResponseStreamEvent>> {
if let Some(data) = line?.strip_prefix("data: ") {
let event = serde_json::from_str(&data)?;
Ok(Some(event))
} else {
Ok(None)
}
}
while let Some(line) = lines.next().await {
if let Some(event) = parse_line(line).transpose() {
let done = event.as_ref().map_or(false, |event| {
event
.choices
.last()
.map_or(false, |choice| choice.finish_reason.is_some())
});
if tx.unbounded_send(event).is_err() {
break;
}
if done {
break;
}
}
}
anyhow::Ok(())
})
.detach();
Ok(rx)
} else {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
#[derive(Deserialize)]
struct OpenAIResponse {
error: OpenAIError,
}
#[derive(Deserialize)]
struct OpenAIError {
message: String,
}
match serde_json::from_str::<OpenAIResponse>(&body) {
Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
"Failed to connect to OpenAI API: {}",
response.error.message,
)),
_ => Err(anyhow!(
"Failed to connect to OpenAI API: {} {}",
response.status(),
body,
)),
}
}
}
#[cfg(test)]
mod tests {
use crate::MessageId;
use super::*;
use fs::FakeFs;
use gpui::{AppContext, TestAppContext};
use project::Project;
fn init_test(cx: &mut TestAppContext) {
cx.foreground().forbid_parking();
cx.update(|cx| {
cx.set_global(SettingsStore::test(cx));
theme::init((), cx);
language::init(cx);
editor::init_settings(cx);
crate::init(cx);
workspace::init_settings(cx);
Project::init_settings(cx);
});
}
#[gpui::test]
async fn test_panel(cx: &mut TestAppContext) {
init_test(cx);
let fs = FakeFs::new(cx.background());
let project = Project::test(fs, [], cx).await;
let (window_id, workspace) = cx.add_window(|cx| Workspace::test_new(project.clone(), cx));
let weak_workspace = workspace.downgrade();
let panel = cx
.spawn(|cx| async move { AssistantPanel::load(weak_workspace, cx).await })
.await
.unwrap();
workspace.update(cx, |workspace, cx| {
workspace.add_panel(panel.clone(), cx);
workspace.toggle_dock(DockPosition::Right, cx);
assert!(workspace.right_dock().read(cx).is_open());
cx.focus(&panel);
});
cx.dispatch_action(window_id, workspace::ToggleZoom);
workspace.read_with(cx, |workspace, cx| {
assert_eq!(workspace.zoomed_view(cx).unwrap(), panel);
})
}
#[gpui::test]
fn test_inserting_and_removing_messages(cx: &mut AppContext) {
let registry = Arc::new(LanguageRegistry::test());
let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
let buffer = conversation.read(cx).buffer.clone();
let message_1 = conversation.read(cx).message_anchors[0].clone();
assert_eq!(
messages(&conversation, cx),
vec![(message_1.id, Role::User, 0..0)]
);
let message_2 = conversation.update(cx, |conversation, cx| {
conversation
.insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)
.unwrap()
});
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..1),
(message_2.id, Role::Assistant, 1..1)
]
);
buffer.update(cx, |buffer, cx| {
buffer.edit([(0..0, "1"), (1..1, "2")], None, cx)
});
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..2),
(message_2.id, Role::Assistant, 2..3)
]
);
let message_3 = conversation.update(cx, |conversation, cx| {
conversation
.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
.unwrap()
});
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..2),
(message_2.id, Role::Assistant, 2..4),
(message_3.id, Role::User, 4..4)
]
);
let message_4 = conversation.update(cx, |conversation, cx| {
conversation
.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
.unwrap()
});
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..2),
(message_2.id, Role::Assistant, 2..4),
(message_4.id, Role::User, 4..5),
(message_3.id, Role::User, 5..5),
]
);
buffer.update(cx, |buffer, cx| {
buffer.edit([(4..4, "C"), (5..5, "D")], None, cx)
});
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..2),
(message_2.id, Role::Assistant, 2..4),
(message_4.id, Role::User, 4..6),
(message_3.id, Role::User, 6..7),
]
);
// Deleting across message boundaries merges the messages.
buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx));
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..3),
(message_3.id, Role::User, 3..4),
]
);
// Undoing the deletion should also undo the merge.
buffer.update(cx, |buffer, cx| buffer.undo(cx));
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..2),
(message_2.id, Role::Assistant, 2..4),
(message_4.id, Role::User, 4..6),
(message_3.id, Role::User, 6..7),
]
);
// Redoing the deletion should also redo the merge.
buffer.update(cx, |buffer, cx| buffer.redo(cx));
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..3),
(message_3.id, Role::User, 3..4),
]
);
// Ensure we can still insert after a merged message.
let message_5 = conversation.update(cx, |conversation, cx| {
conversation
.insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
.unwrap()
});
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..3),
(message_5.id, Role::System, 3..4),
(message_3.id, Role::User, 4..5)
]
);
}
#[gpui::test]
fn test_message_splitting(cx: &mut AppContext) {
let registry = Arc::new(LanguageRegistry::test());
let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
let buffer = conversation.read(cx).buffer.clone();
let message_1 = conversation.read(cx).message_anchors[0].clone();
assert_eq!(
messages(&conversation, cx),
vec![(message_1.id, Role::User, 0..0)]
);
buffer.update(cx, |buffer, cx| {
buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx)
});
let (_, message_2) =
conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
let message_2 = message_2.unwrap();
// We recycle newlines in the middle of a split message
assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_2.id, Role::User, 4..16),
]
);
let (_, message_3) =
conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
let message_3 = message_3.unwrap();
// We don't recycle newlines at the end of a split message
assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_3.id, Role::User, 4..5),
(message_2.id, Role::User, 5..17),
]
);
let (_, message_4) =
conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
let message_4 = message_4.unwrap();
assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_3.id, Role::User, 4..5),
(message_2.id, Role::User, 5..9),
(message_4.id, Role::User, 9..17),
]
);
let (_, message_5) =
conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
let message_5 = message_5.unwrap();
assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_3.id, Role::User, 4..5),
(message_2.id, Role::User, 5..9),
(message_4.id, Role::User, 9..10),
(message_5.id, Role::User, 10..18),
]
);
let (message_6, message_7) = conversation.update(cx, |conversation, cx| {
conversation.split_message(14..16, cx)
});
let message_6 = message_6.unwrap();
let message_7 = message_7.unwrap();
assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_3.id, Role::User, 4..5),
(message_2.id, Role::User, 5..9),
(message_4.id, Role::User, 9..10),
(message_5.id, Role::User, 10..14),
(message_6.id, Role::User, 14..17),
(message_7.id, Role::User, 17..19),
]
);
}
#[gpui::test]
fn test_messages_for_offsets(cx: &mut AppContext) {
let registry = Arc::new(LanguageRegistry::test());
let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
let buffer = conversation.read(cx).buffer.clone();
let message_1 = conversation.read(cx).message_anchors[0].clone();
assert_eq!(
messages(&conversation, cx),
vec![(message_1.id, Role::User, 0..0)]
);
buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));
let message_2 = conversation
.update(cx, |conversation, cx| {
conversation.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx)
})
.unwrap();
buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));
let message_3 = conversation
.update(cx, |conversation, cx| {
conversation.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
})
.unwrap();
buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));
assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");
assert_eq!(
messages(&conversation, cx),
vec![
(message_1.id, Role::User, 0..4),
(message_2.id, Role::User, 4..8),
(message_3.id, Role::User, 8..11)
]
);
assert_eq!(
message_ids_for_offsets(&conversation, &[0, 4, 9], cx),
[message_1.id, message_2.id, message_3.id]
);
assert_eq!(
message_ids_for_offsets(&conversation, &[0, 1, 11], cx),
[message_1.id, message_3.id]
);
fn message_ids_for_offsets(
conversation: &ModelHandle<Conversation>,
offsets: &[usize],
cx: &AppContext,
) -> Vec<MessageId> {
conversation
.read(cx)
.messages_for_offsets(offsets.iter().copied(), cx)
.into_iter()
.map(|message| message.id)
.collect()
}
}
fn messages(
conversation: &ModelHandle<Conversation>,
cx: &AppContext,
) -> Vec<(MessageId, Role, Range<usize>)> {
conversation
.read(cx)
.messages(cx)
.map(|message| (message.id, message.role, message.range))
.collect()
}
}