feat:完善对接dify时返回usage相关信息 (#1715)

alibaba · Feb 3, 2025 · fac2c3e · fac2c3e
1 parent 574d1aa
commit fac2c3e
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/dify.go b/plugins/wasm-go/extensions/ai-proxy/provider/dify.go
@@ -180,6 +180,7 @@ func (d *difyProvider) OnStreamingResponseBody(ctx wrapper.HttpContext, name Api
 func (d *difyProvider) streamResponseDify2OpenAI(ctx wrapper.HttpContext, response *DifyChunkChatResponse) *chatCompletionResponse {
 	var choice chatCompletionChoice
 	var id string
+	var responseUsage usage
 	switch d.config.botType {
 	case BotTypeChat, BotTypeAgent:
 		choice = chatCompletionChoice{
@@ -203,6 +204,13 @@ func (d *difyProvider) streamResponseDify2OpenAI(ctx wrapper.HttpContext, respon
 	}
 	if response.Event == "message_end" || response.Event == "workflow_finished" {
 		choice.FinishReason = finishReasonStop
+		if response.Event == "message_end" {
+			responseUsage = usage{
+				PromptTokens:     response.MetaData.Usage.PromptTokens,
+				CompletionTokens: response.MetaData.Usage.CompletionTokens,
+				TotalTokens:      response.MetaData.Usage.TotalTokens,
+			}
+		}
 	}
 	return &chatCompletionResponse{
 		Id:                id,
@@ -211,6 +219,7 @@ func (d *difyProvider) streamResponseDify2OpenAI(ctx wrapper.HttpContext, respon
 		SystemFingerprint: "",
 		Object:            objectChatCompletionChunk,
 		Choices:           []chatCompletionChoice{choice},
+		Usage:             responseUsage,
 	}
 }