From 5e5524293f43a83aa0f67ef28d6aba4b2790d046 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Wed, 15 Jan 2025 16:35:58 +0800
Subject: [PATCH] Refactor Neo API assistant message handling and function
 integration

- Updated streamChat method to utilize content.Type for function handling, enhancing the interaction with function calls.
- Improved the String method in Content struct to handle function arguments more robustly, including JSON unmarshalling for completed content.
- Set default message type to "text" in NewOpenAI function, ensuring consistent message processing.

These changes enhance the flexibility and maintainability of the Neo API, paving the way for improved assistant functionalities and message management.
---
 neo/assistant/api.go   |  4 ++--
 neo/message/content.go | 31 +++++++++++++++++++++++++++++--
 neo/message/message.go |  1 +
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 7a6ab2c9..b38288fe 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -240,7 +240,7 @@ func (ast *Assistant) streamChat(
 			content.Append(value)
 			if value != "" {
 				// Handle stream
-				res, err := ast.HookStream(c, ctx, messages, content.String(), msg.Type == "tool_calls")
+				res, err := ast.HookStream(c, ctx, messages, content.String(), content.Type == "function")
 				if err == nil && res != nil {
 					if res.Output != "" {
 						value = res.Output
@@ -277,7 +277,7 @@ func (ast *Assistant) streamChat(
 
 				// Call HookDone
 				content.SetStatus(message.ContentStatusDone)
-				res, hookErr := ast.HookDone(c, ctx, messages, content.String(), msg.Type == "tool_calls")
+				res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function")
 				if hookErr == nil && res != nil {
 					if res.Output != "" {
 						chatMessage.New().
diff --git a/neo/message/content.go b/neo/message/content.go
index 972037f3..e79a8f4d 100644
--- a/neo/message/content.go
+++ b/neo/message/content.go
@@ -1,6 +1,8 @@
 package message
 
-import "fmt"
+import (
+	jsoniter "github.com/json-iterator/go"
+)
 
 const (
 	// ContentStatusPending the content status pending
@@ -36,7 +38,32 @@ func NewContent(typ string) *Content {
 // String the content string
 func (c *Content) String() string {
 	if c.Type == "function" {
-		return fmt.Sprintf(`{"id":"%s","type": "function", "function": {"name": "%s", "arguments": "%s"}}`, c.ID, c.Name, c.Bytes)
+
+		var arguments interface{} = string(c.Bytes)
+		if c.Status == ContentStatusDone {
+			var vv interface{} = nil
+			err := jsoniter.Unmarshal(c.Bytes, &vv)
+			if err != nil {
+				return ""
+			}
+			arguments = vv
+		}
+
+		data := map[string]interface{}{
+			"id":   c.ID,
+			"type": "function",
+			"function": map[string]interface{}{
+				"name":      c.Name,
+				"arguments": arguments,
+			},
+		}
+
+		raw, err := jsoniter.MarshalToString(data)
+		if err != nil {
+			return ""
+		}
+
+		return raw
 	}
 	return string(c.Bytes)
 }
diff --git a/neo/message/message.go b/neo/message/message.go
index 72b9c702..97c90ee6 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -97,6 +97,7 @@ func NewOpenAI(data []byte) *Message {
 			return msg
 		}
 
+		msg.Type = "text"
 		if len(message.Choices) > 0 {
 			msg.Text = message.Choices[0].Delta.Content
 		}