From db6ca7a766b8a9d8d6487e5b9c9e25964c687116 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Tue, 5 Nov 2024 02:11:42 +0000
Subject: [PATCH] return first streaming error as http error

---
 llgtrt/src/routes/completions.rs | 12 ++++++++++++
 scripts/req.py                   | 11 ++++++-----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/llgtrt/src/routes/completions.rs b/llgtrt/src/routes/completions.rs
index cc1ee20..2975dda 100644
--- a/llgtrt/src/routes/completions.rs
+++ b/llgtrt/src/routes/completions.rs
@@ -633,11 +633,23 @@ impl ReqInfo {
 async fn completions_stream(
     mut client: ReqInfo,
 ) -> Result<Sse<impl Stream<Item = anyhow::Result<Event>>>, AppError> {
+    let result0 = client
+        .recv
+        .recv()
+        .await
+        .ok_or_else(|| anyhow!("no response"))?;
+    if let Some(err) = result0.response.error {
+        // return as HTTP error, not error event
+        return Err(anyhow!("{}", err).into());
+    }
+
     let response_stream = try_stream! {
         if client.is_run {
             yield client.initial_run();
         }
 
+        yield client.resp_to_event(result0);
+
         while let Some(result) = client.recv.recv().await {
             let is_error = result.response.error.is_some();
             yield client.resp_to_event(result);
diff --git a/scripts/req.py b/scripts/req.py
index f20e875..da8c61d 100644
--- a/scripts/req.py
+++ b/scripts/req.py
@@ -7,7 +7,7 @@
 import time
 import argparse
 
-PROMPT_SIZE = 5
+PROMPT_SIZE = 50_000
 NUM_THREADS = 10
 NUM_REPS = 3
 LLG = False
@@ -167,7 +167,8 @@ def send_one_stream(data: dict) -> list[Results]:
                 break
             data = json.loads(line)
             if data.get("error", None):
-                res.error = json.dumps(data["error"])
+                print("ERROR", data)
+                results[0].error = json.dumps(data["error"])
                 break
 
             if data["object"] == "initial-run":
@@ -283,10 +284,10 @@ def csv_line(lst):
     d = req_data()
     d["n"] = 1
     d["temperature"] = 1.0
-    d["max_tokens"] = 100
-    d["logprobs"] = True
+    d["max_tokens"] = 10
+    # d["logprobs"] = True
     # d["stop"] = ["Xked", "d ask"]
-    if True:
+    if False:
         print(send_one(d))
     else:
         d["llg_log_level"] = "json"