add dify usage

alibaba · Feb 2, 2025 · a1161a2 · a1161a2
1 parent 574d1aa
commit a1161a2
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{
@@ -204,13 +205,21 @@ 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,
 		Created:           time.Now().UnixMilli() / 1000,
 		Model:             ctx.GetStringContext(ctxKeyFinalRequestModel, ""),
 		SystemFingerprint: "",
 		Object:            objectChatCompletionChunk,
 		Choices:           []chatCompletionChoice{choice},
+		Usage:             responseUsage,
 	}
 }