Incrementally diff input coming from GPT

crates/ai/src/ai.rs

@@ -2,27 +2,31 @@ pub mod assistant;
 mod assistant_settings;
 mod refactor;
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 pub use assistant::AssistantPanel;
 use chrono::{DateTime, Local};
 use collections::HashMap;
 use fs::Fs;
-use futures::StreamExt;
-use gpui::AppContext;
+use futures::{io::BufReader, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
+use gpui::{executor::Background, AppContext};
+use isahc::{http::StatusCode, Request, RequestExt};
 use regex::Regex;
 use serde::{Deserialize, Serialize};
 use std::{
     cmp::Reverse,
     ffi::OsStr,
     fmt::{self, Display},
+    io,
     path::PathBuf,
     sync::Arc,
 };
 use util::paths::CONVERSATIONS_DIR;
 
+const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";
+
 // Data types for chat completion requests
 #[derive(Debug, Serialize)]
-struct OpenAIRequest {
+pub struct OpenAIRequest {
     model: String,
     messages: Vec<RequestMessage>,
     stream: bool,
@@ -116,7 +120,7 @@ struct RequestMessage {
 }
 
 #[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
-struct ResponseMessage {
+pub struct ResponseMessage {
     role: Option<Role>,
     content: Option<String>,
 }
@@ -150,7 +154,7 @@ impl Display for Role {
 }
 
 #[derive(Deserialize, Debug)]
-struct OpenAIResponseStreamEvent {
+pub struct OpenAIResponseStreamEvent {
     pub id: Option<String>,
     pub object: String,
     pub created: u32,
@@ -160,14 +164,14 @@ struct OpenAIResponseStreamEvent {
 }
 
 #[derive(Deserialize, Debug)]
-struct Usage {
+pub struct Usage {
     pub prompt_tokens: u32,
     pub completion_tokens: u32,
     pub total_tokens: u32,
 }
 
 #[derive(Deserialize, Debug)]
-struct ChatChoiceDelta {
+pub struct ChatChoiceDelta {
     pub index: u32,
     pub delta: ResponseMessage,
     pub finish_reason: Option<String>,
@@ -190,4 +194,91 @@ struct OpenAIChoice {
 
 pub fn init(cx: &mut AppContext) {
     assistant::init(cx);
+    refactor::init(cx);
+}
+
+pub 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,
+            )),
+        }
+    }
 }