From 29c95c598e380dbe5ff80cd0690a1c4c3770f93d Mon Sep 17 00:00:00 2001
From: creamlike1024 <creamlike1024@outlook.com>
Date: Tue, 13 May 2025 22:01:12 +0800
Subject: [PATCH] cozeChatHelper

---
 relay/channel/coze/dto.go        | 27 ++++++++++++++++++++
 relay/channel/coze/relay-coze.go | 43 +++++++++++++++++++++++++++++---
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/relay/channel/coze/dto.go b/relay/channel/coze/dto.go
index 38fc2f16..4e9afa23 100644
--- a/relay/channel/coze/dto.go
+++ b/relay/channel/coze/dto.go
@@ -49,3 +49,30 @@ type CozeChatUsage struct {
 	OutputCount int `json:"output_count"`
 	InputCount  int `json:"input_count"`
 }
+
+type CozeChatDetailResponse struct {
+	Data   []CozeChatV3MessageDetail `json:"data"`
+	Code   int                       `json:"code"`
+	Msg    string                    `json:"msg"`
+	Detail CozeResponseDetail        `json:"detail"`
+}
+
+type CozeChatV3MessageDetail struct {
+	Id               string          `json:"id"`
+	Role             string          `json:"role"`
+	Type             string          `json:"type"`
+	BotId            string          `json:"bot_id"`
+	ChatId           string          `json:"chat_id"`
+	Content          json.RawMessage `json:"content"`
+	MetaData         json.RawMessage `json:"meta_data"`
+	CreatedAt        int64           `json:"created_at"`
+	SectionId        string          `json:"section_id"`
+	UpdatedAt        int64           `json:"updated_at"`
+	ContentType      string          `json:"content_type"`
+	ConversationId   string          `json:"conversation_id"`
+	ReasoningContent string          `json:"reasoning_content"`
+}
+
+type CozeResponseDetail struct {
+	Logid string `json:"logid"`
+}
diff --git a/relay/channel/coze/relay-coze.go b/relay/channel/coze/relay-coze.go
index 7c16763e..fe630ef6 100644
--- a/relay/channel/coze/relay-coze.go
+++ b/relay/channel/coze/relay-coze.go
@@ -2,12 +2,14 @@ package coze
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"one-api/dto"
 	"one-api/relay/common"
 	relaycommon "one-api/relay/common"
+	"one-api/relay/helper"
 	"one-api/service"
 
 	"github.com/gin-gonic/gin"
@@ -47,14 +49,47 @@ func cozeChatHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rela
 	}
 	// convert coze response to openai response
 	var response dto.TextResponse
-	var cozeResponse CozeChatResponse
+	var cozeResponse CozeChatDetailResponse
+	response.Model = info.UpstreamModelName
 	err = json.Unmarshal(responseBody, &cozeResponse)
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
 	}
-	response.Model = info.UpstreamModelName
-	// TODO: 处理 cozeResponse
-	return nil, nil
+	if cozeResponse.Code != 0 {
+		return service.OpenAIErrorWrapper(errors.New(cozeResponse.Msg), fmt.Sprintf("%d", cozeResponse.Code), http.StatusInternalServerError), nil
+	}
+	// 从上下文获取 usage
+	var usage dto.Usage
+	usage.PromptTokens = c.GetInt("coze_input_count")
+	usage.CompletionTokens = c.GetInt("coze_output_count")
+	usage.TotalTokens = c.GetInt("coze_token_count")
+	response.Usage = usage
+	response.Id = helper.GetResponseID(c)
+
+	var responseContent json.RawMessage
+	for _, data := range cozeResponse.Data {
+		if data.Type == "answer" {
+			responseContent = data.Content
+			response.Created = data.CreatedAt
+		}
+	}
+	// 添加 response.Choices
+	response.Choices = []dto.OpenAITextResponseChoice{
+		{
+			Index:        0,
+			Message:      dto.Message{Role: "assistant", Content: responseContent},
+			FinishReason: "stop",
+		},
+	}
+	jsonResponse, err := json.Marshal(response)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil
+	}
+	c.Writer.Header().Set("Content-Type", "application/json")
+	c.Writer.WriteHeader(resp.StatusCode)
+	_, _ = c.Writer.Write(jsonResponse)
+
+	return nil, &usage
 }
 
 func checkIfChatComplete(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (error, bool) {