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

WIP: Not sure I actually want to rip this out

Browse source
This commit is contained in:
Nathan Sobo 2023-05-26 16:11:58 -06:00
parent 3904971bd8
commit ffbfbe422b
3 changed files with 101 additions and 270 deletions
Download patch file Download diff file Expand all files Collapse all files

128
crates/ai/src/assistant.rs
View file

@ -1,12 +1,14 @@
use crate::{stream_completion, OpenAIRequest, RequestMessage, Role};
use crate::{OpenAIRequest, OpenAIResponseStreamEvent, RequestMessage, Role};
use anyhow::{anyhow, Result};
use editor::{Editor, MultiBuffer};
use futures::StreamExt;
use futures::{io::BufReader, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
use gpui::{
actions, elements::*, Action, AppContext, Entity, ModelHandle, Subscription, Task, View,
ViewContext, ViewHandle, WeakViewHandle, WindowContext,
actions, elements::*, executor::Background, Action, AppContext, Entity, ModelHandle,
Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
};
use isahc::{http::StatusCode, Request, RequestExt};
use language::{language_settings::SoftWrap, Anchor, Buffer};
use std::sync::Arc;
use std::{io, sync::Arc};
use util::{post_inc, ResultExt, TryFutureExt};
use workspace::{
dock::{DockPosition, Panel},
@ -17,8 +19,8 @@ use workspace::{
actions!(assistant, [NewContext, Assist, CancelLastAssist]);
pub fn init(cx: &mut AppContext) {
cx.add_action(ContextEditor::assist);
cx.add_action(ContextEditor::cancel_last_assist);
cx.add_action(Assistant::assist);
cx.capture_action(Assistant::cancel_last_assist);
}
pub enum AssistantPanelEvent {
@ -37,9 +39,7 @@ pub struct AssistantPanel {
impl AssistantPanel {
pub fn new(workspace: &Workspace, cx: &mut ViewContext<Self>) -> Self {
let weak_self = cx.weak_handle();
let pane = cx.add_view(|cx| {
let window_id = cx.window_id();
let mut pane = Pane::new(
workspace.weak_handle(),
workspace.app_state().background_actions,
@ -48,16 +48,15 @@ impl AssistantPanel {
);
pane.set_can_split(false, cx);
pane.set_can_navigate(false, cx);
pane.on_can_drop(move |_, cx| false);
pane.on_can_drop(move |_, _| false);
pane.set_render_tab_bar_buttons(cx, move |pane, cx| {
let this = weak_self.clone();
Flex::row()
.with_child(Pane::render_tab_bar_button(
0,
"icons/plus_12.svg",
Some(("New Context".into(), Some(Box::new(NewContext)))),
cx,
move |_, cx| {},
move |_, _| todo!(),
None,
))
.with_child(Pane::render_tab_bar_button(
@ -123,7 +122,7 @@ impl View for AssistantPanel {
}
impl Panel for AssistantPanel {
fn position(&self, cx: &WindowContext) -> DockPosition {
fn position(&self, _: &WindowContext) -> DockPosition {
DockPosition::Right
}
@ -131,9 +130,11 @@ impl Panel for AssistantPanel {
matches!(position, DockPosition::Right)
}
fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext<Self>) {}
fn set_position(&mut self, _: DockPosition, _: &mut ViewContext<Self>) {
// TODO!
}
fn size(&self, cx: &WindowContext) -> f32 {
fn size(&self, _: &WindowContext) -> f32 {
self.width.unwrap_or(480.)
}
@ -164,7 +165,7 @@ impl Panel for AssistantPanel {
if let Some(workspace) = this.workspace.upgrade(cx) {
workspace.update(cx, |workspace, cx| {
let focus = this.pane.read(cx).has_focus();
let editor = Box::new(cx.add_view(|cx| ContextEditor::new(cx)));
let editor = Box::new(cx.add_view(|cx| Assistant::new(cx)));
Pane::add_item(workspace, &this.pane, editor, true, focus, None, cx);
})
}
@ -180,7 +181,8 @@ impl Panel for AssistantPanel {
("Assistant Panel".into(), None)
}
fn should_change_position_on_event(event: &Self::Event) -> bool {
fn should_change_position_on_event(_: &Self::Event) -> bool {
// TODO!
false
}
@ -201,7 +203,7 @@ impl Panel for AssistantPanel {
}
}
struct ContextEditor {
struct Assistant {
messages: Vec<Message>,
editor: ViewHandle<Editor>,
completion_count: usize,
@ -210,10 +212,10 @@ struct ContextEditor {
struct PendingCompletion {
id: usize,
task: Task<Option<()>>,
_task: Task<Option<()>>,
}
impl ContextEditor {
impl Assistant {
fn new(cx: &mut ViewContext<Self>) -> Self {
let messages = vec![Message {
role: Role::User,
@ -264,15 +266,26 @@ impl ContextEditor {
if let Some(api_key) = std::env::var("OPENAI_API_KEY").log_err() {
let stream = stream_completion(api_key, cx.background_executor().clone(), request);
let content = cx.add_model(|cx| Buffer::new(0, "", cx));
let response_buffer = cx.add_model(|cx| Buffer::new(0, "", cx));
self.messages.push(Message {
role: Role::Assistant,
content: content.clone(),
content: response_buffer.clone(),
});
let next_request_buffer = cx.add_model(|cx| Buffer::new(0, "", cx));
self.messages.push(Message {
role: Role::User,
content: next_request_buffer.clone(),
});
self.editor.update(cx, |editor, cx| {
editor.buffer().update(cx, |multibuffer, cx| {
multibuffer.push_excerpts_with_context_lines(
content.clone(),
response_buffer.clone(),
vec![Anchor::MIN..Anchor::MAX],
0,
cx,
);
multibuffer.push_excerpts_with_context_lines(
next_request_buffer,
vec![Anchor::MIN..Anchor::MAX],
0,
cx,
@ -286,7 +299,7 @@ impl ContextEditor {
while let Some(message) = messages.next().await {
let mut message = message?;
if let Some(choice) = message.choices.pop() {
content.update(&mut cx, |content, cx| {
response_buffer.update(&mut cx, |content, cx| {
let text: Arc<str> = choice.delta.content?.into();
content.edit([(content.len()..content.len(), text)], None, cx);
Some(())
@ -307,23 +320,23 @@ impl ContextEditor {
self.pending_completions.push(PendingCompletion {
id: post_inc(&mut self.completion_count),
task,
_task: task,
});
}
}
fn cancel_last_assist(&mut self, _: &CancelLastAssist, cx: &mut ViewContext<Self>) {
fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
if self.pending_completions.pop().is_none() {
cx.propagate_action();
}
}
}
impl Entity for ContextEditor {
impl Entity for Assistant {
type Event = ();
}
impl View for ContextEditor {
impl View for Assistant {
fn ui_name() -> &'static str {
"ContextEditor"
}
@ -338,7 +351,7 @@ impl View for ContextEditor {
}
}
impl Item for ContextEditor {
impl Item for Assistant {
fn tab_content<V: View>(
&self,
_: Option<usize>,
@ -353,3 +366,60 @@ struct Message {
role: Role,
content: ModelHandle<Buffer>,
}
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("https://api.openai.com/v1/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() {
tx.unbounded_send(event).log_err();
}
}
anyhow::Ok(())
})
.detach();
Ok(rx)
} else {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
Err(anyhow!(
"Failed to connect to OpenAI API: {} {}",
response.status(),
body,
))
}
}