git-sync/crawlab
1
0
Fork 0
mirror of https://github.com/crawlab-team/crawlab.git synced 2026-01-21 17:21:09 +01:00
Code Issues Packages Projects Releases Wiki Activity

refactor: update LLMResponseUsage struct and ChatMessage usage aggregation

Browse Source 
- Renamed fields in LLMResponseUsage struct to better reflect their purpose, changing InputTokens and OutputTokens to PromptTokens and CompletionTokens.
- Modified GetUsage method in ChatMessage to accept an optional contents parameter, improving flexibility in usage tracking.
- Ensured consistent field naming and improved clarity in token usage aggregation logic.
This commit is contained in:
Marvin Zhang
2025-04-14 16:29:32 +08:00
parent 8273507d20
commit 67b51fdc83
2 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
core/entity/llm.go
View File

@@ -1,7 +1,7 @@
package entity
type LLMResponseUsage struct {
InputTokens int `json:"input_tokens" bson:"input_tokens"`
OutputTokens int `json:"output_tokens" bson:"output_tokens"`
TotalTokens int `json:"total_tokens" bson:"total_tokens"`
PromptTokens int `json:"prompt_tokens" bson:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens" bson:"completion_tokens"`
TotalTokens int `json:"total_tokens" bson:"total_tokens"`
}

13
core/models/models/chat_message.go
View File

@@ -57,16 +57,19 @@ func (m *ChatMessage) GetContent() string {
return result
}
func (m *ChatMessage) GetUsage() *entity.LLMResponseUsage {
if len(m.Contents) == 0 {
func (m *ChatMessage) GetUsage(contents []ChatMessageContent) *entity.LLMResponseUsage {
if contents == nil {
contents = m.Contents
}
if len(contents) == 0 {
return nil
}
var usage entity.LLMResponseUsage
for _, content := range m.Contents {
for _, content := range contents {
if content.Usage != nil {
// Accumulate usage
usage.InputTokens += content.Usage.InputTokens
usage.OutputTokens += content.Usage.OutputTokens
usage.PromptTokens += content.Usage.PromptTokens
usage.CompletionTokens += content.Usage.CompletionTokens
usage.TotalTokens += content.Usage.TotalTokens
}
}