diff --git a/core/entity/llm.go b/core/entity/llm.go
index 87302b32..d2fcb822 100644
--- a/core/entity/llm.go
+++ b/core/entity/llm.go
@@ -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"`
 }
diff --git a/core/models/models/chat_message.go b/core/models/models/chat_message.go
index 6253ef5b..d196a916 100644
--- a/core/models/models/chat_message.go
+++ b/core/models/models/chat_message.go
@@ -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
 		}
 	}