From fac2c3e7a3ef924de381a8453c63b18a3c918b75 Mon Sep 17 00:00:00 2001
From: kai2321 <70182551+kai2321@users.noreply.github.com>
Date: Mon, 3 Feb 2025 08:35:00 +0800
Subject: [PATCH] =?UTF-8?q?feat:=E5=AE=8C=E5=96=84=E5=AF=B9=E6=8E=A5dify?=
 =?UTF-8?q?=E6=97=B6=E8=BF=94=E5=9B=9Eusage=E7=9B=B8=E5=85=B3=E4=BF=A1?=
 =?UTF-8?q?=E6=81=AF=20(#1715)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 plugins/wasm-go/extensions/ai-proxy/provider/dify.go | 9 +++++++++
 1 file changed, 9 insertions(+)

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,
 	}
 }