From 6ac0406ad4c4110ebd4c07d7054b43ccc1ef6692 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Sat, 1 Feb 2025 11:26:11 +0800
Subject: [PATCH] Improve OpenAI message parsing and error handling

- Enhanced error handling in OpenAI message processing
- Added colored console output for JSON parsing errors
- Updated Error struct to support flexible code type
- Improved message type and done state setting for various error scenarios
- Added support for parsing different error message formats
---
 neo/message/message.go | 35 ++++++++++++++++++++++++++++++++---
 openai/types.go        |  2 +-
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/neo/message/message.go b/neo/message/message.go
index 3b8d3dba..bbcbe1c1 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -191,7 +191,11 @@ func NewOpenAI(data []byte) *Message {
 	case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"tool_calls"`):
 		var toolCalls openai.ToolCalls
 		if err := jsoniter.Unmarshal(data, &toolCalls); err != nil {
-			msg.Text = err.Error() + "\n" + string(data)
+			color.Red("JSON parse error: %s", err.Error())
+			color.White(string(data))
+			msg.Text = "JSON parse error\n" + string(data)
+			msg.Type = "error"
+			msg.IsDone = true
 			return msg
 		}
 
@@ -205,7 +209,11 @@ func NewOpenAI(data []byte) *Message {
 	case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`):
 		var message openai.Message
 		if err := jsoniter.Unmarshal(data, &message); err != nil {
-			msg.Text = err.Error() + "\n" + string(data)
+			color.Red("JSON parse error: %s", err.Error())
+			color.White(string(data))
+			msg.Text = "JSON parse error\n" + string(data)
+			msg.Type = "error"
+			msg.IsDone = true
 			return msg
 		}
 
@@ -214,14 +222,35 @@ func NewOpenAI(data []byte) *Message {
 			msg.Text = message.Choices[0].Delta.Content
 		}
 
+	case strings.Index(text, `{"code":`) == 0:
+		var errorMessage openai.Error
+		if err := jsoniter.UnmarshalFromString(text, &errorMessage); err != nil {
+			color.Red("JSON parse error: %s", err.Error())
+			color.White(string(data))
+			msg.Text = "JSON parse error\n" + string(data)
+			msg.Type = "error"
+			msg.IsDone = true
+			return msg
+		}
+		msg.Type = "error"
+		msg.Text = errorMessage.Message
+		msg.IsDone = true
+		break
+
 	case strings.Contains(text, `{"error":{`):
 		var errorMessage openai.ErrorMessage
 		if err := jsoniter.Unmarshal(data, &errorMessage); err != nil {
-			msg.Text = err.Error() + "\n" + string(data)
+			color.Red("JSON parse error: %s", err.Error())
+			color.White(string(data))
+			msg.Text = "JSON parse error\n" + string(data)
+			msg.Type = "error"
+			msg.IsDone = true
 			return msg
 		}
 		msg.Type = "error"
 		msg.Text = errorMessage.Error.Message
+		msg.IsDone = true
+		break
 
 	case strings.Contains(text, `[DONE]`):
 		msg.IsDone = true
diff --git a/openai/types.go b/openai/types.go
index 3c99ce7a..0c6d58b4 100644
--- a/openai/types.go
+++ b/openai/types.go
@@ -48,5 +48,5 @@ type Error struct {
 	Message string      `json:"message,omitempty"`
 	Type    string      `json:"type,omitempty"`
 	Param   interface{} `json:"param,omitempty"`
-	Code    string      `json:"code,omitempty"`
+	Code    any         `json:"code,omitempty"` // string or int
 }