From a45a3d6241b28522683189a4587135f18aaacaaa Mon Sep 17 00:00:00 2001
From: "Laisky.Cai" <github@laisky.com>
Date: Wed, 19 Feb 2025 08:10:04 +0000
Subject: [PATCH] feat: enhance reasoning token handling in OpenAI adaptor

---
 relay/adaptor/openai/main.go | 9 +++++++++
 relay/controller/helper.go   | 3 +++
 2 files changed, 12 insertions(+)

diff --git a/relay/adaptor/openai/main.go b/relay/adaptor/openai/main.go
index 5065a5b1..a6ef3084 100644
--- a/relay/adaptor/openai/main.go
+++ b/relay/adaptor/openai/main.go
@@ -66,6 +66,9 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E
 				if choice.Delta.Reasoning != nil {
 					reasoningText += *choice.Delta.Reasoning
 				}
+				if choice.Delta.ReasoningContent != nil {
+					reasoningText += *choice.Delta.ReasoningContent
+				}
 
 				responseText += conv.AsString(choice.Delta.Content)
 			}
@@ -99,6 +102,12 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E
 		return ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "", nil
 	}
 
+	// If there is no reasoning tokens in the completion, we should count the reasoning tokens in the response.
+	if len(reasoningText) > 0 &&
+		(usage.CompletionTokensDetails == nil || usage.CompletionTokensDetails.ReasoningTokens == 0) {
+		usage.CompletionTokens += CountToken(reasoningText)
+	}
+
 	return nil, reasoningText + responseText, usage
 }
 
diff --git a/relay/controller/helper.go b/relay/controller/helper.go
index 63ac8e96..021f703c 100644
--- a/relay/controller/helper.go
+++ b/relay/controller/helper.go
@@ -104,6 +104,9 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.M
 
 	completionRatio := billingratio.GetCompletionRatio(textRequest.Model, meta.ChannelType)
 	promptTokens := usage.PromptTokens
+	// It appears that DeepSeek's official service automatically merges ReasoningTokens into CompletionTokens,
+	// but the behavior of third-party providers may differ, so for now we do not add them manually.
+	// completionTokens := usage.CompletionTokens + usage.CompletionTokensDetails.ReasoningTokens
 	completionTokens := usage.CompletionTokens
 	quota = int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio))
 	if ratio != 0 && quota <= 0 {