feat: 完善stream_options

README.md

@@ -72,7 +72,8 @@
 
 ## 比原版One API多出的配置
 - `STREAMING_TIMEOUT`：设置流式一次回复的超时时间，默认为 30 秒
-
+- `DIFY_DEBUG`：设置 Dify 渠道是否输出工作流和节点信息到客户端，默认为 `true`， 可选值为 `true` 和 `false`
+- `FORCE_STREAM_OPTION`：覆盖客户端stream_options参数，请求上游返回流模式usage，目前仅支持 `OpenAI` 渠道类型
 ## 部署
 ### 部署要求
 - 本地数据库（默认）：SQLite（Docker 部署默认使用 SQLite，必须挂载 `/data` 目录到宿主机）

common/env.go

@@ -24,3 +24,15 @@ func GetEnvOrDefaultString(env string, defaultValue string) string {
 	}
 	return os.Getenv(env)
 }
+
+func GetEnvOrDefaultBool(env string, defaultValue bool) bool {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	b, err := strconv.ParseBool(os.Getenv(env))
+	if err != nil {
+		SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %t", env, err.Error(), defaultValue))
+		return defaultValue
+	}
+	return b
+}

constant/env.go

@@ -5,3 +5,7 @@ import (
 )
 
 var StreamingTimeout = common.GetEnvOrDefault("STREAMING_TIMEOUT", 30)
+var DifyDebug = common.GetEnvOrDefaultBool("DIFY_DEBUG", true)
+
+// ForceStreamOption 覆盖请求参数，强制返回usage信息
+var ForceStreamOption = common.GetEnvOrDefaultBool("FORCE_STREAM_OPTION", true)

dto/text_request.go

@@ -11,6 +11,7 @@ type GeneralOpenAIRequest struct {
 	Messages         []Message       `json:"messages,omitempty"`
 	Prompt           any             `json:"prompt,omitempty"`
 	Stream           bool            `json:"stream,omitempty"`
+	StreamOptions    *StreamOptions  `json:"stream_options,omitempty"`
 	MaxTokens        uint            `json:"max_tokens,omitempty"`
 	Temperature      float64         `json:"temperature,omitempty"`
 	TopP             float64         `json:"top_p,omitempty"`
@@ -43,6 +44,10 @@ type OpenAIFunction struct {
 	Parameters  any    `json:"parameters,omitempty"`
 }
 
+type StreamOptions struct {
+	IncludeUsage bool `json:"include_usage,omitempty"`
+}
+
 func (r GeneralOpenAIRequest) GetMaxTokens() int64 {
 	return int64(r.MaxTokens)
 }

dto/text_response.go

@@ -102,10 +102,12 @@ type ChatCompletionsStreamResponse struct {
 	Model             string                                `json:"model"`
 	SystemFingerprint *string                               `json:"system_fingerprint"`
 	Choices           []ChatCompletionsStreamResponseChoice `json:"choices"`
+	Usage             *Usage                                `json:"usage"`
 }
 
 type ChatCompletionsStreamResponseSimple struct {
 	Choices []ChatCompletionsStreamResponseChoice `json:"choices"`
+	Usage   *Usage                                `json:"usage"`
 }
 
 type CompletionsStreamResponse struct {

relay/channel/aws/adaptor.go

@@ -68,7 +68,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
 
 func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
-		err, usage = awsStreamHandler(c, info, a.RequestMode)
+		err, usage = awsStreamHandler(c, resp, info, a.RequestMode)
 	} else {
 		err, usage = awsHandler(c, info, a.RequestMode)
 	}

relay/channel/aws/relay-aws.go

@@ -13,6 +13,7 @@ import (
 	relaymodel "one-api/dto"
 	"one-api/relay/channel/claude"
 	relaycommon "one-api/relay/common"
+	"one-api/service"
 	"strings"
 	"time"
 
@@ -112,7 +113,7 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode int) (*
 	return nil, &usage
 }
 
-func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode int) (*relaymodel.OpenAIErrorWithStatusCode, *relaymodel.Usage) {
+func awsStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo, requestMode int) (*relaymodel.OpenAIErrorWithStatusCode, *relaymodel.Usage) {
 	awsCli, err := newAwsClient(c, info)
 	if err != nil {
 		return wrapErr(errors.Wrap(err, "newAwsClient")), nil
@@ -162,7 +163,6 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode i
 	c.Stream(func(w io.Writer) bool {
 		event, ok := <-stream.Events()
 		if !ok {
-			c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
 			return false
 		}
 
@@ -214,6 +214,17 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode i
 			return false
 		}
 	})
-
+	if info.ShouldIncludeUsage {
+		response := service.GenerateFinalUsageResponse(id, createdTime, info.UpstreamModelName, usage)
+		err := service.ObjectData(c, response)
+		if err != nil {
+			common.SysError("send final response failed: " + err.Error())
+		}
+	}
+	service.Done(c)
+	err = resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapperLocal(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
 	return nil, &usage
 }

relay/channel/claude/relay-claude.go

@@ -330,22 +330,15 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 			response.Created = createdTime
 			response.Model = info.UpstreamModelName
 
-			jsonStr, err := json.Marshal(response)
+			err = service.ObjectData(c, response)
 			if err != nil {
-				common.SysError("error marshalling stream response: " + err.Error())
-				return true
+				common.SysError(err.Error())
 			}
-			c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)})
 			return true
 		case <-stopChan:
-			c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
 			return false
 		}
 	})
-	err := resp.Body.Close()
-	if err != nil {
-		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
-	}
 	if requestMode == RequestModeCompletion {
 		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
 	} else {
@@ -356,6 +349,18 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 			usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, usage.PromptTokens)
 		}
 	}
+	if info.ShouldIncludeUsage {
+		response := service.GenerateFinalUsageResponse(responseId, createdTime, info.UpstreamModelName, *usage)
+		err := service.ObjectData(c, response)
+		if err != nil {
+			common.SysError("send final response failed: " + err.Error())
+		}
+	}
+	service.Done(c)
+	err := resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapperLocal(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
 	return nil, usage
 }
 

relay/channel/dify/relay-dify.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"one-api/common"
+	"one-api/constant"
 	"one-api/dto"
 	relaycommon "one-api/relay/common"
 	"one-api/service"
@@ -48,9 +49,9 @@ func streamResponseDify2OpenAI(difyResponse DifyChunkChatCompletionResponse) *dt
 		Model:   "dify",
 	}
 	var choice dto.ChatCompletionsStreamResponseChoice
-	if difyResponse.Event == "workflow_started" {
+	if constant.DifyDebug && difyResponse.Event == "workflow_started" {
 		choice.Delta.SetContentString("Workflow: " + difyResponse.Data.WorkflowId + "\n")
-	} else if difyResponse.Event == "node_started" {
+	} else if constant.DifyDebug && difyResponse.Event == "node_started" {
 		choice.Delta.SetContentString("Node: " + difyResponse.Data.NodeId + "\n")
 	} else if difyResponse.Event == "message" {
 		choice.Delta.SetContentString(difyResponse.Answer)

relay/channel/ollama/adaptor.go

@@ -59,8 +59,10 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
 func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		var responseText string
-		err, responseText, _ = openai.OpenaiStreamHandler(c, resp, info)
-		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+		err, usage, responseText, _ = openai.OpenaiStreamHandler(c, resp, info)
+		if usage == nil || usage.TotalTokens == 0 || (usage.PromptTokens+usage.CompletionTokens) == 0 {
+			usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+		}
 	} else {
 		if info.RelayMode == relayconstant.RelayModeEmbeddings {
 			err, usage = ollamaEmbeddingHandler(c, resp, info.PromptTokens, info.UpstreamModelName, info.RelayMode)

relay/channel/openai/adaptor.go

@@ -89,9 +89,11 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
 	if info.IsStream {
 		var responseText string
 		var toolCount int
-		err, responseText, toolCount = OpenaiStreamHandler(c, resp, info)
-		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
-		usage.CompletionTokens += toolCount * 7
+		err, usage, responseText, toolCount = OpenaiStreamHandler(c, resp, info)
+		if usage == nil || usage.TotalTokens == 0 || (usage.PromptTokens+usage.CompletionTokens) == 0 {
+			usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+			usage.CompletionTokens += toolCount * 7
+		}
 	} else {
 		err, usage = OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
 	}

relay/channel/openai/relay-openai.go

@@ -18,9 +18,10 @@ import (
 	"time"
 )
 
-func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, string, int) {
+func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, *dto.Usage, string, int) {
 	//checkSensitive := constant.ShouldCheckCompletionSensitive()
 	var responseTextBuilder strings.Builder
+	var usage dto.Usage
 	toolCount := 0
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
@@ -62,17 +63,26 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 				streamItems = append(streamItems, data)
 			}
 		}
+		// 计算token
 		streamResp := "[" + strings.Join(streamItems, ",") + "]"
 		switch info.RelayMode {
 		case relayconstant.RelayModeChatCompletions:
 			var streamResponses []dto.ChatCompletionsStreamResponseSimple
 			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
 			if err != nil {
+				// 一次性解析失败，逐个解析
 				common.SysError("error unmarshalling stream response: " + err.Error())
 				for _, item := range streamItems {
 					var streamResponse dto.ChatCompletionsStreamResponseSimple
 					err := json.Unmarshal(common.StringToByteSlice(item), &streamResponse)
 					if err == nil {
+						if streamResponse.Usage != nil {
+							if streamResponse.Usage.TotalTokens != 0 {
+								usage.PromptTokens += streamResponse.Usage.PromptTokens
+								usage.CompletionTokens += streamResponse.Usage.CompletionTokens
+								usage.TotalTokens += streamResponse.Usage.TotalTokens
+							}
+						}
 						for _, choice := range streamResponse.Choices {
 							responseTextBuilder.WriteString(choice.Delta.GetContentString())
 							if choice.Delta.ToolCalls != nil {
@@ -89,6 +99,13 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 				}
 			} else {
 				for _, streamResponse := range streamResponses {
+					if streamResponse.Usage != nil {
+						if streamResponse.Usage.TotalTokens != 0 {
+							usage.PromptTokens += streamResponse.Usage.PromptTokens
+							usage.CompletionTokens += streamResponse.Usage.CompletionTokens
+							usage.TotalTokens += streamResponse.Usage.TotalTokens
+						}
+					}
 					for _, choice := range streamResponse.Choices {
 						responseTextBuilder.WriteString(choice.Delta.GetContentString())
 						if choice.Delta.ToolCalls != nil {
@@ -107,6 +124,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 			var streamResponses []dto.CompletionsStreamResponse
 			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
 			if err != nil {
+				// 一次性解析失败，逐个解析
 				common.SysError("error unmarshalling stream response: " + err.Error())
 				for _, item := range streamItems {
 					var streamResponse dto.CompletionsStreamResponse
@@ -133,7 +151,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 	}()
 	service.SetEventStreamHeaders(c)
 	isFirst := true
-	ticker := time.NewTicker(time.Duration(constant.StreamingTimeout))
+	ticker := time.NewTicker(time.Duration(constant.StreamingTimeout) * time.Second)
 	defer ticker.Stop()
 	c.Stream(func(w io.Writer) bool {
 		select {
@@ -145,7 +163,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 				isFirst = false
 				info.FirstResponseTime = time.Now()
 			}
-			ticker.Reset(time.Duration(constant.StreamingTimeout))
+			ticker.Reset(time.Duration(constant.StreamingTimeout) * time.Second)
 			if strings.HasPrefix(data, "data: [DONE]") {
 				data = data[:12]
 			}
@@ -159,10 +177,10 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 	})
 	err := resp.Body.Close()
 	if err != nil {
-		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "", toolCount
+		return service.OpenAIErrorWrapperLocal(err, "close_response_body_failed", http.StatusInternalServerError), nil, "", toolCount
 	}
 	wg.Wait()
-	return nil, responseTextBuilder.String(), toolCount
+	return nil, &usage, responseTextBuilder.String(), toolCount
 }
 
 func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {

relay/channel/perplexity/adaptor.go

@@ -55,8 +55,10 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
 func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
 	if info.IsStream {
 		var responseText string
-		err, responseText, _ = openai.OpenaiStreamHandler(c, resp, info)
-		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+		err, usage, responseText, _ = openai.OpenaiStreamHandler(c, resp, info)
+		if usage == nil || usage.TotalTokens == 0 || (usage.PromptTokens+usage.CompletionTokens) == 0 {
+			usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+		}
 	} else {
 		err, usage = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
 	}

relay/channel/zhipu_4v/adaptor.go

@@ -57,9 +57,11 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
 	if info.IsStream {
 		var responseText string
 		var toolCount int
-		err, responseText, toolCount = openai.OpenaiStreamHandler(c, resp, info)
-		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
-		usage.CompletionTokens += toolCount * 7
+		err, usage, responseText, toolCount = openai.OpenaiStreamHandler(c, resp, info)
+		if usage == nil || usage.TotalTokens == 0 || (usage.PromptTokens+usage.CompletionTokens) == 0 {
+			usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
+			usage.CompletionTokens += toolCount * 7
+		}
 	} else {
 		err, usage = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
 	}

relay/common/relay_info.go

@@ -9,24 +9,26 @@ import (
 )
 
 type RelayInfo struct {
-	ChannelType       int
-	ChannelId         int
-	TokenId           int
-	UserId            int
-	Group             string
-	TokenUnlimited    bool
-	StartTime         time.Time
-	FirstResponseTime time.Time
-	ApiType           int
-	IsStream          bool
-	RelayMode         int
-	UpstreamModelName string
-	RequestURLPath    string
-	ApiVersion        string
-	PromptTokens      int
-	ApiKey            string
-	Organization      string
-	BaseUrl           string
+	ChannelType          int
+	ChannelId            int
+	TokenId              int
+	UserId               int
+	Group                string
+	TokenUnlimited       bool
+	StartTime            time.Time
+	FirstResponseTime    time.Time
+	ApiType              int
+	IsStream             bool
+	RelayMode            int
+	UpstreamModelName    string
+	RequestURLPath       string
+	ApiVersion           string
+	PromptTokens         int
+	ApiKey               string
+	Organization         string
+	BaseUrl              string
+	SupportStreamOptions bool
+	ShouldIncludeUsage   bool
 }
 
 func GenRelayInfo(c *gin.Context) *RelayInfo {
@@ -65,6 +67,9 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 	if info.ChannelType == common.ChannelTypeAzure {
 		info.ApiVersion = GetAPIVersion(c)
 	}
+	if info.ChannelType == common.ChannelTypeOpenAI || info.ChannelType == common.ChannelTypeAnthropic || info.ChannelType == common.ChannelTypeAws {
+		info.SupportStreamOptions = true
+	}
 	return info
 }
 

relay/relay-text.go

@@ -77,7 +77,7 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 
 	// map model name
 	modelMapping := c.GetString("model_mapping")
-	isModelMapped := false
+	//isModelMapped := false
 	if modelMapping != "" && modelMapping != "{}" {
 		modelMap := make(map[string]string)
 		err := json.Unmarshal([]byte(modelMapping), &modelMap)
@@ -87,7 +87,7 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 		if modelMap[textRequest.Model] != "" {
 			textRequest.Model = modelMap[textRequest.Model]
 			// set upstream model name
-			isModelMapped = true
+			//isModelMapped = true
 		}
 	}
 	relayInfo.UpstreamModelName = textRequest.Model
@@ -130,33 +130,38 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 		return openaiErr
 	}
 
+	// 如果不支持StreamOptions，将StreamOptions设置为nil
+	if !relayInfo.SupportStreamOptions || !textRequest.Stream {
+		textRequest.StreamOptions = nil
+	} else {
+		// 如果支持StreamOptions，且请求中没有设置StreamOptions，根据配置文件设置StreamOptions
+		if constant.ForceStreamOption {
+			textRequest.StreamOptions = &dto.StreamOptions{
+				IncludeUsage: true,
+			}
+		}
+	}
+
+	if textRequest.StreamOptions != nil && textRequest.StreamOptions.IncludeUsage {
+		relayInfo.ShouldIncludeUsage = textRequest.StreamOptions.IncludeUsage
+	}
+
 	adaptor := GetAdaptor(relayInfo.ApiType)
 	if adaptor == nil {
 		return service.OpenAIErrorWrapperLocal(fmt.Errorf("invalid api type: %d", relayInfo.ApiType), "invalid_api_type", http.StatusBadRequest)
 	}
 	adaptor.Init(relayInfo, *textRequest)
 	var requestBody io.Reader
-	if relayInfo.ApiType == relayconstant.APITypeOpenAI {
-		if isModelMapped {
-			jsonStr, err := json.Marshal(textRequest)
-			if err != nil {
-				return service.OpenAIErrorWrapperLocal(err, "marshal_text_request_failed", http.StatusInternalServerError)
-			}
-			requestBody = bytes.NewBuffer(jsonStr)
-		} else {
-			requestBody = c.Request.Body
-		}
-	} else {
-		convertedRequest, err := adaptor.ConvertRequest(c, relayInfo.RelayMode, textRequest)
-		if err != nil {
-			return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError)
-		}
-		jsonData, err := json.Marshal(convertedRequest)
-		if err != nil {
-			return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError)
-		}
-		requestBody = bytes.NewBuffer(jsonData)
+
+	convertedRequest, err := adaptor.ConvertRequest(c, relayInfo.RelayMode, textRequest)
+	if err != nil {
+		return service.OpenAIErrorWrapperLocal(err, "convert_request_failed", http.StatusInternalServerError)
 	}
+	jsonData, err := json.Marshal(convertedRequest)
+	if err != nil {
+		return service.OpenAIErrorWrapperLocal(err, "json_marshal_failed", http.StatusInternalServerError)
+	}
+	requestBody = bytes.NewBuffer(jsonData)
 
 	statusCodeMappingStr := c.GetString("status_code_mapping")
 	resp, err := adaptor.DoRequest(c, relayInfo, requestBody)

service/usage_helpr.go

@@ -24,3 +24,15 @@ func ResponseText2Usage(responseText string, modeName string, promptTokens int)
 	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 	return usage, err
 }
+
+func GenerateFinalUsageResponse(id string, createAt int64, model string, usage dto.Usage) *dto.ChatCompletionsStreamResponse {
+	return &dto.ChatCompletionsStreamResponse{
+		Id:                id,
+		Object:            "chat.completion.chunk",
+		Created:           createAt,
+		Model:             model,
+		SystemFingerprint: nil,
+		Choices:           nil,
+		Usage:             &usage,
+	}
+}