use crate::inline_prompt_editor::CodegenStatus;use futures::{SinkExt, StreamExt, channel::mpsc};use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Task};use language_model::{ConfiguredModel, LanguageModelRegistry, LanguageModelRequest};use language_models::provider::anthropic::telemetry::{ AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType,};use std::time::Instant;use terminal::Terminal;use uuid::Uuid; pub struct TerminalCodegen { pub status: CodegenStatus, terminal: Entity<Terminal>, generation: Task<()>, pub message_id: Option<String>, transaction: Option<TerminalTransaction>, session_id: Uuid,} impl EventEmitter<CodegenEvent> for TerminalCodegen {} impl TerminalCodegen { pub fn new(terminal: Entity<Terminal>, session_id: Uuid) -> Self { Self { terminal, status: CodegenStatus::Idle, generation: Task::ready(()), message_id: None, transaction: None, session_id, } } pub fn session_id(&self) -> Uuid { self.session_id } pub fn start(&mut self, prompt_task: Task<LanguageModelRequest>, cx: &mut Context<Self>) { let Some(ConfiguredModel { model, .. }) = LanguageModelRegistry::read_global(cx).inline_assistant_model() else { return; }; let anthropic_reporter = AnthropicEventReporter::new(&model, cx); let session_id = self.session_id; let model_telemetry_id = model.telemetry_id(); let model_provider_id = model.provider_id().to_string(); self.status = CodegenStatus::Pending; self.transaction = Some(TerminalTransaction::start(self.terminal.clone())); self.generation = cx.spawn(async move |this, cx| { let prompt = prompt_task.await; let response = model.stream_completion_text(prompt, cx).await; let generate = async { let message_id = response .as_ref() .ok() .and_then(|response| response.message_id.clone()); let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1); let task = cx.background_spawn({ let message_id = message_id.clone(); let anthropic_reporter = anthropic_reporter.clone(); async move { let mut response_latency = None; let request_start = Instant::now(); let task = async { let mut chunks = response?.stream; while let Some(chunk) = chunks.next().await { if response_latency.is_none() { response_latency = Some(request_start.elapsed()); } let chunk = chunk?; hunks_tx.send(chunk).await?; } anyhow::Ok(()) }; let result = task.await; let error_message = result.as_ref().err().map(|error| error.to_string()); telemetry::event!( "Assistant Responded", session_id = session_id.to_string(), kind = "inline_terminal", phase = "response", model = model_telemetry_id, model_provider = model_provider_id, language_name = Option::<&str>::None, message_id = message_id, response_latency = response_latency, error_message = error_message, ); anthropic_reporter.report(AnthropicEventData { completion_type: AnthropicCompletionType::Terminal, event: AnthropicEventType::Response, language_name: None, message_id, }); result?; anyhow::Ok(()) } }); this.update(cx, |this, _| { this.message_id = message_id; })?; while let Some(hunk) = hunks_rx.next().await { this.update(cx, |this, cx| { if let Some(transaction) = &mut this.transaction { transaction.push(hunk, cx); cx.notify(); } })?; } task.await?; anyhow::Ok(()) }; let result = generate.await; this.update(cx, |this, cx| { if let Err(error) = result { this.status = CodegenStatus::Error(error); } else { this.status = CodegenStatus::Done; } cx.emit(CodegenEvent::Finished); cx.notify(); }) .ok(); }); cx.notify(); } pub fn completion(&self) -> Option<String> { self.transaction .as_ref() .map(|transaction| transaction.completion.clone()) } pub fn stop(&mut self, cx: &mut Context<Self>) { self.status = CodegenStatus::Done; self.generation = Task::ready(()); cx.emit(CodegenEvent::Finished); cx.notify(); } pub fn complete(&mut self, cx: &mut Context<Self>) { if let Some(transaction) = self.transaction.take() { transaction.complete(cx); } } pub fn undo(&mut self, cx: &mut Context<Self>) { if let Some(transaction) = self.transaction.take() { transaction.undo(cx); } }} #[derive(Copy, Clone, Debug)]pub enum CodegenEvent { Finished,} #[cfg(not(target_os = "windows"))]pub const CLEAR_INPUT: &str = "\x15";#[cfg(target_os = "windows")]pub const CLEAR_INPUT: &str = "\x03";const CARRIAGE_RETURN: &str = "\x0d"; struct TerminalTransaction { completion: String, terminal: Entity<Terminal>,} impl TerminalTransaction { pub fn start(terminal: Entity<Terminal>) -> Self { Self { completion: String::new(), terminal, } } pub fn push(&mut self, hunk: String, cx: &mut App) { // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal let input = Self::sanitize_input(hunk); self.completion.push_str(&input); self.terminal .update(cx, |terminal, _| terminal.input(input.into_bytes())); } pub fn undo(self, cx: &mut App) { self.terminal .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.as_bytes())); } pub fn complete(self, cx: &mut App) { self.terminal .update(cx, |terminal, _| terminal.input(CARRIAGE_RETURN.as_bytes())); } fn sanitize_input(mut input: String) -> String { input.retain(|c| c != '\r' && c != '\n'); input }}
use crate::inline_prompt_editor::CodegenStatus;
use futures::{SinkExt, StreamExt, channel::mpsc};
use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Task};
use language_model::{ConfiguredModel, LanguageModelRegistry, LanguageModelRequest};
use language_models::provider::anthropic::telemetry::{
AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType,
};
use std::time::Instant;
use terminal::Terminal;
use uuid::Uuid;
pub struct TerminalCodegen {
pub status: CodegenStatus,
terminal: Entity<Terminal>,
generation: Task<()>,
pub message_id: Option<String>,
transaction: Option<TerminalTransaction>,
session_id: Uuid,
}
impl EventEmitter<CodegenEvent> for TerminalCodegen {}
impl TerminalCodegen {
pub fn new(terminal: Entity<Terminal>, session_id: Uuid) -> Self {
Self {
terminal,
status: CodegenStatus::Idle,
generation: Task::ready(()),
message_id: None,
transaction: None,
session_id,
}
}
pub fn session_id(&self) -> Uuid {
self.session_id
}
pub fn start(&mut self, prompt_task: Task<LanguageModelRequest>, cx: &mut Context<Self>) {
let Some(ConfiguredModel { model, .. }) =
LanguageModelRegistry::read_global(cx).inline_assistant_model()
else {
return;
};
let anthropic_reporter = AnthropicEventReporter::new(&model, cx);
let session_id = self.session_id;
let model_telemetry_id = model.telemetry_id();
let model_provider_id = model.provider_id().to_string();
self.status = CodegenStatus::Pending;
self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
self.generation = cx.spawn(async move |this, cx| {
let prompt = prompt_task.await;
let response = model.stream_completion_text(prompt, cx).await;
let generate = async {
let message_id = response
.as_ref()
.ok()
.and_then(|response| response.message_id.clone());
let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
let task = cx.background_spawn({
let message_id = message_id.clone();
let anthropic_reporter = anthropic_reporter.clone();
async move {
let mut response_latency = None;
let request_start = Instant::now();
let task = async {
let mut chunks = response?.stream;
while let Some(chunk) = chunks.next().await {
if response_latency.is_none() {
response_latency = Some(request_start.elapsed());
}
let chunk = chunk?;
hunks_tx.send(chunk).await?;
}
anyhow::Ok(())
};
let result = task.await;
let error_message = result.as_ref().err().map(|error| error.to_string());
telemetry::event!(
"Assistant Responded",
session_id = session_id.to_string(),
kind = "inline_terminal",
phase = "response",
model = model_telemetry_id,
model_provider = model_provider_id,
language_name = Option::<&str>::None,
message_id = message_id,
response_latency = response_latency,
error_message = error_message,
);
anthropic_reporter.report(AnthropicEventData {
completion_type: AnthropicCompletionType::Terminal,
event: AnthropicEventType::Response,
language_name: None,
message_id,
});
result?;
anyhow::Ok(())
}
});
this.update(cx, |this, _| {
this.message_id = message_id;
})?;
while let Some(hunk) = hunks_rx.next().await {
this.update(cx, |this, cx| {
if let Some(transaction) = &mut this.transaction {
transaction.push(hunk, cx);
cx.notify();
}
})?;
}
task.await?;
anyhow::Ok(())
};
let result = generate.await;
this.update(cx, |this, cx| {
if let Err(error) = result {
this.status = CodegenStatus::Error(error);
} else {
this.status = CodegenStatus::Done;
}
cx.emit(CodegenEvent::Finished);
cx.notify();
})
.ok();
});
cx.notify();
}
pub fn completion(&self) -> Option<String> {
self.transaction
.as_ref()
.map(|transaction| transaction.completion.clone())
}
pub fn stop(&mut self, cx: &mut Context<Self>) {
self.status = CodegenStatus::Done;
self.generation = Task::ready(());
cx.emit(CodegenEvent::Finished);
cx.notify();
}
pub fn complete(&mut self, cx: &mut Context<Self>) {
if let Some(transaction) = self.transaction.take() {
transaction.complete(cx);
}
}
pub fn undo(&mut self, cx: &mut Context<Self>) {
if let Some(transaction) = self.transaction.take() {
transaction.undo(cx);
}
}
}
#[derive(Copy, Clone, Debug)]
pub enum CodegenEvent {
Finished,
}
#[cfg(not(target_os = "windows"))]
pub const CLEAR_INPUT: &str = "\x15";
#[cfg(target_os = "windows")]
pub const CLEAR_INPUT: &str = "\x03";
const CARRIAGE_RETURN: &str = "\x0d";
struct TerminalTransaction {
completion: String,
terminal: Entity<Terminal>,
}
impl TerminalTransaction {
pub fn start(terminal: Entity<Terminal>) -> Self {
Self {
completion: String::new(),
terminal,
}
}
pub fn push(&mut self, hunk: String, cx: &mut App) {
// Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
let input = Self::sanitize_input(hunk);
self.completion.push_str(&input);
self.terminal
.update(cx, |terminal, _| terminal.input(input.into_bytes()));
}
pub fn undo(self, cx: &mut App) {
self.terminal
.update(cx, |terminal, _| terminal.input(CLEAR_INPUT.as_bytes()));
}
pub fn complete(self, cx: &mut App) {
self.terminal
.update(cx, |terminal, _| terminal.input(CARRIAGE_RETURN.as_bytes()));
}
fn sanitize_input(mut input: String) -> String {
input.retain(|c| c != '\r' && c != '\n');
input
}
}