From 8d76fc77fb770915177d2412def509e3ef829e60 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 1 Nov 2023 13:24:50 +0800
Subject: refactor: use reqwest-eventsource as sse client

---
 src/client/openai.rs | 48 ++++++++++++++++++++++++++++++------------------
 1 file changed, 30 insertions(+), 18 deletions(-)

(limited to 'src/client/openai.rs')

diff --git a/src/client/openai.rs b/src/client/openai.rs
index 80969e4..4ae3c5f 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -5,9 +5,9 @@ use crate::repl::ReplyStreamHandler;
 
 use anyhow::{anyhow, bail, Result};
 use async_trait::async_trait;
-use eventsource_stream::Eventsource;
 use futures_util::StreamExt;
 use reqwest::{Client as ReqwestClient, RequestBuilder};
+use reqwest_eventsource::{Error as EventSourceError, Event, RequestBuilderExt};
 use serde::Deserialize;
 use serde_json::{json, Value};
 use std::env;
@@ -104,23 +104,35 @@ pub async fn openai_send_message_streaming(
     builder: RequestBuilder,
     handler: &mut ReplyStreamHandler,
 ) -> Result<()> {
-    let res = builder.send().await?;
-    if !res.status().is_success() {
-        let data: Value = res.json().await?;
-        if let Some(err_msg) = data["error"]["message"].as_str() {
-            bail!("{err_msg}");
-        }
-        bail!("Request failed");
-    }
-    let mut stream = res.bytes_stream().eventsource();
-    while let Some(part) = stream.next().await {
-        let chunk = part?.data;
-        if chunk == "[DONE]" {
-            break;
-        }
-        let data: Value = serde_json::from_str(&chunk)?;
-        if let Some(text) = data["choices"][0]["delta"]["content"].as_str() {
-            handler.text(text)?;
+    let mut es = builder.eventsource()?;
+    while let Some(event) = es.next().await {
+        match event {
+            Ok(Event::Open) => {}
+            Ok(Event::Message(message)) => {
+                if message.data == "[DONE]" {
+                    break;
+                }
+                let data: Value = serde_json::from_str(&message.data)?;
+                if let Some(text) = data["choices"][0]["delta"]["content"].as_str() {
+                    handler.text(text)?;
+                }
+            }
+            Err(err) => {
+                match err {
+                    EventSourceError::InvalidStatusCode(_, res) => {
+                        let data: Value = res.json().await?;
+                        if let Some(err_msg) = data["error"]["message"].as_str() {
+                            bail!("{err_msg}");
+                        }
+                        bail!("Request failed");
+                    }
+                    EventSourceError::StreamEnded => {}
+                    _ => {
+                        bail!("{}", err);
+                    }
+                }
+                es.close();
+            }
         }
     }
 
-- 
cgit v1.2.3