From 17081393fce54e450311d7798af3a16af6a86316 Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Thu, 4 Apr 2024 23:32:28 +0800
Subject: [PATCH] fix: fix usage & ratio

---
 common/model-ratio.go       | 1 +
 relay/channel/zhipu/main.go | 6 +++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index aa75042e..d8356dc2 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -91,6 +91,7 @@ var ModelRatio = map[string]float64{
 	"glm-4":                     0.1 * RMB,
 	"glm-4v":                    0.1 * RMB,
 	"glm-3-turbo":               0.005 * RMB,
+	"embedding-2":               0.0005 * RMB,
 	"chatglm_turbo":             0.3572, // ￥0.005 / 1k tokens
 	"chatglm_pro":               0.7143, // ￥0.01 / 1k tokens
 	"chatglm_std":               0.3572, // ￥0.005 / 1k tokens
diff --git a/relay/channel/zhipu/main.go b/relay/channel/zhipu/main.go
index 4f4bdf48..f54e0504 100644
--- a/relay/channel/zhipu/main.go
+++ b/relay/channel/zhipu/main.go
@@ -285,7 +285,11 @@ func embeddingResponseZhipu2OpenAI(response *EmbeddingRespone) *openai.Embedding
 		Object: "list",
 		Data:   make([]openai.EmbeddingResponseItem, 0, len(response.Embeddings)),
 		Model:  response.Model,
-		Usage:  model.Usage{TotalTokens: response.Usage.TotalTokens},
+		Usage: model.Usage{
+			PromptTokens:     response.PromptTokens,
+			CompletionTokens: response.CompletionTokens,
+			TotalTokens:      response.Usage.TotalTokens,
+		},
 	}
 
 	for _, item := range response.Embeddings {