姚章浩
31 lines
1.1 KiB
Go
31 lines
1.1 KiB
Go
package agent
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
)
|
|
|
|
func TestRecordTokenUsageAggregatesAgentRounds(t *testing.T) {
|
|
engine := &AgentEngine{}
|
|
engine.recordTokenUsage(types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120, CachedTokens: 40})
|
|
engine.recordTokenUsage(types.TokenUsage{PromptTokens: 140, CompletionTokens: 30, TotalTokens: 170, CachedTokens: 60})
|
|
|
|
wantTotal := types.TokenUsage{PromptTokens: 240, CompletionTokens: 50, TotalTokens: 290, CachedTokens: 100}
|
|
if engine.totalUsage != wantTotal {
|
|
t.Fatalf("totalUsage = %#v, want %#v", engine.totalUsage, wantTotal)
|
|
}
|
|
wantLast := types.TokenUsage{PromptTokens: 140, CompletionTokens: 30, TotalTokens: 170, CachedTokens: 60}
|
|
if engine.lastUsage != wantLast {
|
|
t.Fatalf("lastUsage = %#v, want %#v", engine.lastUsage, wantLast)
|
|
}
|
|
}
|
|
|
|
func TestRecordTokenUsageIgnoresEmptyUsage(t *testing.T) {
|
|
engine := &AgentEngine{}
|
|
engine.recordTokenUsage(types.TokenUsage{})
|
|
if engine.totalUsage != (types.TokenUsage{}) || engine.lastUsage != (types.TokenUsage{}) {
|
|
t.Fatalf("empty usage changed state: total=%#v last=%#v", engine.totalUsage, engine.lastUsage)
|
|
}
|
|
}
|