From 92cd46d64f2713a52624dd04ef8aa4567cc53edf Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Sun, 10 Nov 2024 00:36:08 +0800
Subject: [PATCH] feat: able to use ENFORCE_INCLUDE_USAGE to enforce include
 usage in response

---
 README.md                | 1 +
 common/config/config.go  | 2 ++
 relay/controller/text.go | 3 ++-
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index f27a6794..fb137c23 100644
--- a/README.md
+++ b/README.md
@@ -400,6 +400,7 @@ graph LR
 26. `METRIC_SUCCESS_RATE_THRESHOLD`：请求成功率阈值，默认为 `0.8`。
 27. `INITIAL_ROOT_TOKEN`：如果设置了该值，则在系统首次启动时会自动创建一个值为该环境变量值的 root 用户令牌。
 28. `INITIAL_ROOT_ACCESS_TOKEN`：如果设置了该值，则在系统首次启动时会自动创建一个值为该环境变量的 root 用户创建系统管理令牌。
+29. `ENFORCE_INCLUDE_USAGE`：是否强制在 stream 模型下返回 usage，默认不开启，可选值为 `true` 和 `false`。
 
 ### 命令行参数
 1. `--port <port_number>`: 指定服务器监听的端口号，默认为 `3000`。
diff --git a/common/config/config.go b/common/config/config.go
index 43f56862..2eb894ef 100644
--- a/common/config/config.go
+++ b/common/config/config.go
@@ -160,3 +160,5 @@ var OnlyOneLogFile = env.Bool("ONLY_ONE_LOG_FILE", false)
 var RelayProxy = env.String("RELAY_PROXY", "")
 var UserContentRequestProxy = env.String("USER_CONTENT_REQUEST_PROXY", "")
 var UserContentRequestTimeout = env.Int("USER_CONTENT_REQUEST_TIMEOUT", 30)
+
+var EnforceIncludeUsage = env.Bool("ENFORCE_INCLUDE_USAGE", false)
diff --git a/relay/controller/text.go b/relay/controller/text.go
index 52ee9949..57f98812 100644
--- a/relay/controller/text.go
+++ b/relay/controller/text.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/songquanpeng/one-api/common/config"
 	"io"
 	"net/http"
 
@@ -84,7 +85,7 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode {
 }
 
 func getRequestBody(c *gin.Context, meta *meta.Meta, textRequest *model.GeneralOpenAIRequest, adaptor adaptor.Adaptor) (io.Reader, error) {
-	if meta.APIType == apitype.OpenAI && meta.OriginModelName == meta.ActualModelName && meta.ChannelType != channeltype.Baichuan {
+	if !config.EnforceIncludeUsage && meta.APIType == apitype.OpenAI && meta.OriginModelName == meta.ActualModelName && meta.ChannelType != channeltype.Baichuan {
 		// no need to convert request for openai
 		return c.Request.Body, nil
 	}