From 040c48b9b392e3329d3f1de8eeb1b8773129773c Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 24 Apr 2024 12:38:59 +0800
Subject: refactor: handling of response error (#433)

---
 src/client/cohere.rs | 45 ++++++++++++++++++++++++---------------------
 1 file changed, 24 insertions(+), 21 deletions(-)

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

diff --git a/src/client/cohere.rs b/src/client/cohere.rs
index db2ad28..4d713b6 100644
--- a/src/client/cohere.rs
+++ b/src/client/cohere.rs
@@ -78,7 +78,7 @@ pub(crate) async fn send_message(builder: RequestBuilder) -> Result<String> {
     let status = res.status();
     let data: Value = res.json().await?;
     if status != 200 {
-        check_error(&data)?;
+        catch_error(&data, status.as_u16())?;
     }
     let output = extract_text(&data)?;
     Ok(output.to_string())
@@ -89,9 +89,10 @@ pub(crate) async fn send_message_streaming(
     handler: &mut ReplyHandler,
 ) -> Result<()> {
     let res = builder.send().await?;
-    if res.status() != 200 {
+    let status = res.status();
+    if status != 200 {
         let data: Value = res.json().await?;
-        check_error(&data)?;
+        catch_error(&data, status.as_u16())?;
     } else {
         let handle = |value: &str| -> Result<()> {
             let value: Value = serde_json::from_str(value)?;
@@ -105,24 +106,7 @@ pub(crate) async fn send_message_streaming(
     Ok(())
 }
 
-fn extract_text(data: &Value) -> Result<&str> {
-    match data["text"].as_str() {
-        Some(text) => Ok(text),
-        None => {
-            bail!("Invalid response data: {data}")
-        }
-    }
-}
-
-fn check_error(data: &Value) -> Result<()> {
-    if let Some(message) = data["message"].as_str() {
-        bail!("{message}");
-    } else {
-        bail!("Error {}", data);
-    }
-}
-
-pub(crate) fn build_body(data: SendData, model: &Model) -> Result<Value> {
+fn build_body(data: SendData, model: &Model) -> Result<Value> {
     let SendData {
         mut messages,
         temperature,
@@ -195,3 +179,22 @@ pub(crate) fn build_body(data: SendData, model: &Model) -> Result<Value> {
 
     Ok(body)
 }
+
+fn catch_error(data: &Value, status: u16) -> Result<()> {
+    debug!("Invalid response, status: {status}, data: {data}");
+
+    if let Some(message) = data["message"].as_str() {
+        bail!("{message}");
+    } else {
+        bail!("Invalid response, status: {status}, data: {data}");
+    }
+}
+
+fn extract_text(data: &Value) -> Result<&str> {
+    match data["text"].as_str() {
+        Some(text) => Ok(text),
+        None => {
+            bail!("Invalid response data: {data}")
+        }
+    }
+}
-- 
cgit v1.2.3