diff --git a/plugins/wasm-go/extensions/ai-proxy/provider/dify.go b/plugins/wasm-go/extensions/ai-proxy/provider/dify.go index 28b6dec79..5c21cc3b6 100644 --- 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, } }