feat(ai-proxy): add chat completion usage for dashscope director (#6389)

* feat(ai-proxy): add chat completion usage for dashscope director * tidy go.sum * stream resp only have usage field at last chunk
erda-project · Jul 8, 2024 · 6f783df · 6f783df
1 parent de03578
commit 6f783df
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 3 deletions.
diff --git a/go.mod b/go.mod
@@ -130,7 +130,7 @@ require (
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/russross/blackfriday/v2 v2.1.0
 	github.com/sabhiram/go-gitignore v0.0.0-20201211210132-54b8a0bf510f
-	github.com/sashabaranov/go-openai v1.17.9
+	github.com/sashabaranov/go-openai v1.26.2
 	github.com/satori/go.uuid v1.2.0
 	github.com/scylladb/gocqlx v1.5.0
 	github.com/shirou/gopsutil v3.21.11+incompatible

diff --git a/go.sum b/go.sum
@@ -2469,8 +2469,8 @@ github.com/sanity-io/litter v1.2.0/go.mod h1:JF6pZUFgu2Q0sBZ+HSV35P8TVPI1TTzEwyu
 github.com/santhosh-tekuri/jsonschema v1.2.4 h1:hNhW8e7t+H1vgY+1QeEQpveR6D4+OwKPXCfD2aieJis=
 github.com/santhosh-tekuri/jsonschema v1.2.4/go.mod h1:TEAUOeZSmIxTTuHatJzrvARHiuO9LYd+cIxzgEHCQI4=
 github.com/sasha-s/go-deadlock v0.2.0/go.mod h1:StQn567HiB1fF2yJ44N9au7wOhrPS3iZqiDbRupzT10=
-github.com/sashabaranov/go-openai v1.17.9 h1:QEoBiGKWW68W79YIfXWEFZ7l5cEgZBV4/Ow3uy+5hNY=
-github.com/sashabaranov/go-openai v1.17.9/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
+github.com/sashabaranov/go-openai v1.26.2 h1:cVlQa3gn3eYqNXRW03pPlpy6zLG52EU4g0FrWXc0EFI=
+github.com/sashabaranov/go-openai v1.26.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/satori/go.uuid v1.2.0 h1:0uYX9dsZ2yD7q2RtLRtPSdGDWzjeM3TbMJP9utgA0ww=
 github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0=
 github.com/scylladb/go-reflectx v1.0.1 h1:b917wZM7189pZdlND9PbIJ6NQxfDPfBvUaQ7cjj1iZQ=

diff --git a/internal/apps/ai-proxy/filters/dashscope-director/sdk/converter.go b/internal/apps/ai-proxy/filters/dashscope-director/sdk/converter.go
@@ -97,5 +97,12 @@ func ConvertDsStreamChunkToOpenAIFormat(dsChunk DsRespStreamChunk, modelName str
 		Model:   modelName,
 		Choices: ocs,
 	}
+	if dsChunk.Usage != nil {
+		openaiChunk.Usage = &openai.Usage{
+			PromptTokens:     int(dsChunk.Usage.InputTokens),
+			CompletionTokens: int(dsChunk.Usage.OutputTokens),
+			TotalTokens:      int(dsChunk.Usage.TotalTokens),
+		}
+	}
 	return &openaiChunk, nil
 }
diff --git a/internal/apps/ai-proxy/filters/dashscope-director/sdk/resp.go b/internal/apps/ai-proxy/filters/dashscope-director/sdk/resp.go
@@ -18,6 +18,7 @@ type (
 	DsRespStreamChunk struct {
 		Output    DsRespStreamChunkOutput `json:"output,omitempty"`
 		RequestID string                  `json:"request_id,omitempty"`
+		Usage     *DsRespStreamUsage      `json:"usage,omitempty"`
 	}
 	DsRespStreamChunkOutput struct {
 		Choices []DsRespStreamChunkOutputChoice `json:"choices,omitempty"`
@@ -34,4 +35,9 @@ type (
 	DsRespStreamChunkOutputChoiceMessagePart struct {
 		Text string `json:"text,omitempty"`
 	}
+	DsRespStreamUsage struct {
+		TotalTokens  uint64 `json:"total_tokens"`
+		InputTokens  uint64 `json:"input_tokens"`
+		OutputTokens uint64 `json:"output_tokens"`
+	}
 )