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

refactor: update LLMResponseUsage struct and ChatMessage model for BSON compatibility

Browse Source 
- Added BSON tags to fields in LLMResponseUsage struct to ensure proper serialization with MongoDB.
- Updated the Usage field in ChatMessage model to exclude it from JSON serialization while retaining BSON compatibility.
- Introduced GetUsage method in ChatMessage to aggregate token usage from contents, enhancing usage tracking capabilities.
This commit is contained in:
Marvin Zhang
2025-04-12 21:25:57 +08:00
parent 2549abdc61
commit 8273507d20
2 changed files with 20 additions and 4 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"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
InputTokens int `json:"input_tokens" bson:"input_tokens"`
OutputTokens int `json:"output_tokens" bson:"output_tokens"`
TotalTokens int `json:"total_tokens" bson:"total_tokens"`
}

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

@@ -17,7 +17,7 @@ type ChatMessage struct {
Model string `json:"model" bson:"model" description:"AI model used"`
Status string `json:"status" bson:"status" description:"Message status (pending/completed/failed)"`
Error string `json:"error,omitempty" bson:"error,omitempty" description:"Error message if failed"`
Usage *entity.LLMResponseUsage `json:"usage,omitempty" bson:"usage,omitempty" description:"Usage"`
Usage *entity.LLMResponseUsage `json:"usage,omitempty" bson:"-" description:"Usage"`
}
func (m *ChatMessage) GetContent() string {
@@ -56,3 +56,19 @@ func (m *ChatMessage) GetContent() string {
return result
}
func (m *ChatMessage) GetUsage() *entity.LLMResponseUsage {
if len(m.Contents) == 0 {
return nil
}
var usage entity.LLMResponseUsage
for _, content := range m.Contents {
if content.Usage != nil {
// Accumulate usage
usage.InputTokens += content.Usage.InputTokens
usage.OutputTokens += content.Usage.OutputTokens
usage.TotalTokens += content.Usage.TotalTokens
}
}
return &usage
}