3978622050
- Refactored AI gateway to utilize new cost management structures for usage tracking. - Replaced deprecated token extraction methods with a unified cost parsing approach. - Enhanced usage fallback mechanisms and introduced detailed usage metrics in responses. - Added new metering functionality to record AI usage and costs effectively. - Updated tests to reflect changes in usage parsing and cost calculations. - Introduced new API endpoints for retrieving AI usage summaries and pricing information.
73 lines
2.0 KiB
Go
73 lines
2.0 KiB
Go
package cost
|
|
|
|
import "testing"
|
|
|
|
func TestParseUsageOpenAI(t *testing.T) {
|
|
payload := []byte(`{
|
|
"usage": {
|
|
"prompt_tokens": 1200,
|
|
"completion_tokens": 340,
|
|
"total_tokens": 1540,
|
|
"prompt_tokens_details": {"cached_tokens": 800},
|
|
"completion_tokens_details": {"reasoning_tokens": 50}
|
|
}
|
|
}`)
|
|
u := ParseUsage(payload)
|
|
if u.PromptTokens != 1200 || u.CompletionTokens != 340 {
|
|
t.Fatalf("unexpected tokens: %+v", u)
|
|
}
|
|
if u.CachedInputTokens != 800 || u.ReasoningTokens != 50 {
|
|
t.Fatalf("unexpected details: %+v", u)
|
|
}
|
|
}
|
|
|
|
func TestParseUsageFallback(t *testing.T) {
|
|
u := ParseUsage([]byte(`{"choices":[]}`))
|
|
if u.TotalTokens != 1 {
|
|
t.Fatalf("expected fallback 1, got %d", u.TotalTokens)
|
|
}
|
|
}
|
|
|
|
func TestComputeCostMicroEUR(t *testing.T) {
|
|
price := ModelPrice{
|
|
InputMicroEURPerMTok: 1_000_000,
|
|
OutputMicroEURPerMTok: 2_000_000,
|
|
}
|
|
usage := UsageDetail{
|
|
PromptTokens: 1000,
|
|
CompletionTokens: 500,
|
|
CachedInputTokens: 200,
|
|
TotalTokens: 1500,
|
|
}
|
|
cost, estimated := ComputeCostMicroEUR(usage, price, true)
|
|
if estimated {
|
|
t.Fatal("should not be estimated when price found")
|
|
}
|
|
// uncached 800 * 1 + cached 200 * 0.5 + output 500 * 2 = 800+100+1000 = 1900 micro (cached uses half input)
|
|
if cost < 1800 || cost > 2000 {
|
|
t.Fatalf("unexpected cost %d", cost)
|
|
}
|
|
}
|
|
|
|
func TestComputeCostUnknownModel(t *testing.T) {
|
|
usage := UsageDetail{TotalTokens: 1_000_000}
|
|
cost, estimated := ComputeCostMicroEUR(usage, ModelPrice{}, false)
|
|
if !estimated {
|
|
t.Fatal("expected estimated")
|
|
}
|
|
if cost != fallbackInputMicroEURPerMTok {
|
|
t.Fatalf("expected fallback cost %d, got %d", fallbackInputMicroEURPerMTok, cost)
|
|
}
|
|
}
|
|
|
|
func TestMergeStreamUsage(t *testing.T) {
|
|
acc := UsageDetail{}
|
|
chunk1 := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":0,"total_tokens":10}}`)
|
|
chunk2 := []byte(`{"usage":{"prompt_tokens":100,"completion_tokens":40,"total_tokens":140}}`)
|
|
acc = MergeStreamUsage(acc, chunk1)
|
|
acc = MergeStreamUsage(acc, chunk2)
|
|
if acc.TotalTokens != 140 {
|
|
t.Fatalf("expected 140, got %d", acc.TotalTokens)
|
|
}
|
|
}
|