fix: fix json_util

yiwen-ai · Nov 27, 2023 · 98f4b60 · 98f4b60
1 parent 25e396e
commit 98f4b60
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 33 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "jarvis"
-version = "1.1.6"
+version = "1.1.7"
 edition = "2021"
 rust-version = "1.64"
 description = ""

diff --git a/src/json_util.rs b/src/json_util.rs
@@ -50,7 +50,7 @@ impl RawJSONArray {
 
     fn skip_space(&mut self) {
         while self.offset < self.chars.len() {
-            if self.chars[self.offset].is_whitespace() {
+            if self.chars[self.offset].is_whitespace() || self.chars[self.offset].is_control() {
                 self.offset += 1;
             } else {
                 break;
@@ -61,7 +61,7 @@ impl RawJSONArray {
     fn skip_space_v(&self) -> usize {
         let mut offset = self.offset;
         while offset < self.chars.len() {
-            if self.chars[offset].is_whitespace() {
+            if self.chars[offset].is_whitespace() || self.chars[self.offset].is_control() {
                 offset += 1;
             } else {
                 break;
@@ -234,7 +234,10 @@ impl RawJSONArray {
                     self.result.push(']');
                 }
                 _ => {
-                    self.result.push(self.chars[self.offset]);
+                    let c = self.chars[self.offset];
+                    if !c.is_control() {
+                        self.result.push(c);
+                    }
                     self.offset += 1;
                 }
             }
@@ -273,8 +276,8 @@ mod tests {
                 err: None,
             },
             Case {
-                input: r#"[" "]"#.to_string(),
-                output: r#"[" "]"#.to_string(),
+                input: r#"[" ❤️‍🔥🧑‍🤝‍🧑"]"#.to_string(),
+                output: r#"[" ❤️‍🔥🧑‍🤝‍🧑"]"#.to_string(),
                 err: None,
             },
             Case {
@@ -330,7 +333,7 @@ mod tests {
                 err: None,
             },
             Case {
-                input: r#"["\ "]"#.to_string(),
+                input: r#"["\ "]"#.to_string(), // with a control char
                 output: r#"["\\ "]"#.to_string(),
                 err: None,
             },

diff --git a/src/openai.rs b/src/openai.rs
@@ -10,7 +10,7 @@ use reqwest::{header, Client, ClientBuilder, Identity, Response};
 use serde::{de::DeserializeOwned, Serialize};
 use std::{path::Path, str::FromStr, string::ToString};
 use tiktoken_rs::{num_tokens_from_messages, ChatCompletionRequestMessage};
-use tokio::time::Duration;
+use tokio::time::{sleep, Duration};
 
 use crate::conf::AI;
 use crate::json_util::RawJSONArray;
@@ -61,8 +61,8 @@ impl AIModel {
     // return (recommend, high)
     pub fn translating_segment_tokens(&self) -> (usize, usize) {
         match self {
-            AIModel::GPT3_5 => (3000, 3400),
-            AIModel::GPT4 => (3000, 3400),
+            AIModel::GPT3_5 => (3000, 3200),
+            AIModel::GPT4 => (3000, 3200),
         }
     }
 
@@ -260,19 +260,12 @@ impl OpenAI {
             total_tokens: 0,
         });
 
-        let real_tokens_rate: f32 = if usage.prompt_tokens > 1000 {
-            usage.completion_tokens as f32 / (usage.prompt_tokens as f32 - 90f32)
-        } else {
-            1.0f32
-        };
-
         let elapsed = ctx.start.elapsed().as_millis() as u32;
         ctx.set_kvs(vec![
             ("elapsed", elapsed.into()),
             ("prompt_tokens", usage.prompt_tokens.into()),
             ("completion_tokens", usage.completion_tokens.into()),
             ("total_tokens", usage.total_tokens.into()),
-            ("real_tokens_rate", real_tokens_rate.into()),
             ("speed", (usage.total_tokens * 1000 / elapsed).into()),
         ])
         .await;
@@ -523,6 +516,7 @@ impl OpenAI {
         match Self::check_chat_response(res) {
             Ok(rt) => Ok(rt),
             Err(err) if err.code == 429 || err.code > 500 => {
+                sleep(Duration::from_secs(3)).await;
                 ctx.set("retry_because", err.to_string().into()).await;
                 rand_index += 1;
                 (api_url, headers) = self.get_params(&model_name, rand_index);
@@ -616,6 +610,7 @@ impl OpenAI {
         match Self::check_chat_response(res) {
             Ok(rt) => Ok(rt),
             Err(err) if err.code == 429 || err.code > 500 => {
+                sleep(Duration::from_secs(3)).await;
                 ctx.set("retry_because", err.to_string().into()).await;
                 rand_index += 1;
                 (api_url, headers) = self.get_params(&model_name, rand_index);
@@ -654,11 +649,7 @@ impl OpenAI {
                             return Err(HTTPError {
                                 code: 452,
                                 message: "Content was triggered the filtering model".to_string(),
-                                data: choice
-                                    .message
-                                    .content
-                                    .clone()
-                                    .map(serde_json::Value::String),
+                                data: serde_json::to_value(rt).ok(),
                             });
                         }
 
@@ -667,23 +658,15 @@ impl OpenAI {
                                 code: 422,
                                 message: "Incomplete output due to max_tokens parameter"
                                     .to_string(),
-                                data: choice
-                                    .message
-                                    .content
-                                    .clone()
-                                    .map(serde_json::Value::String),
+                                data: serde_json::to_value(rt).ok(),
                             })
                         }
 
                         reason => {
                             return Err(HTTPError {
                                 code: 500,
                                 message: format!("Unknown finish reason: {}", reason),
-                                data: choice
-                                    .message
-                                    .content
-                                    .clone()
-                                    .map(serde_json::Value::String),
+                                data: serde_json::to_value(rt).ok(),
                             });
                         }
                     }
@@ -752,6 +735,7 @@ impl OpenAI {
         match Self::check_chat_response(res) {
             Ok(rt) => Ok(rt),
             Err(err) if err.code == 429 || err.code > 500 => {
+                sleep(Duration::from_secs(3)).await;
                 ctx.set("retry_because", err.to_string().into()).await;
                 rand_index += 1;
                 (api_url, headers) = self.get_params(&model_name, rand_index);