lyric.tngl.sh / mistermorph
Unified Agent + reusable Go agent core.
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

fix: handle prompt cache usage details

Lyric e0557c06 17bee38c

+211 -6
+134 -4
providers/uniai/client.go
··· 48 48 type Client struct { 49 49 provider string 50 50 model string 51 + pricing *uniaiapi.PricingCatalog 51 52 requestTimeout time.Duration 52 53 temperature *float64 53 54 reasoningEffort string ··· 59 60 60 61 func New(cfg Config) *Client { 61 62 provider := strings.ToLower(strings.TrimSpace(cfg.Provider)) 63 + pricing := cfg.Pricing 64 + if pricing == nil { 65 + pricing = uniaiapi.DefaultPricingCatalog() 66 + } 62 67 63 68 openAIBase := normalizeOpenAIBase(cfg.Endpoint) 64 69 openAIKey := strings.TrimSpace(cfg.APIKey) ··· 79 84 OpenAIAPIBase: openAIBase, 80 85 OpenAIModel: strings.TrimSpace(cfg.Model), 81 86 ChatHeaders: cloneStringMap(cfg.Headers), 82 - Pricing: cfg.Pricing, 83 87 AzureOpenAIAPIKey: strings.TrimSpace(azureAPIKey), 84 88 AzureOpenAIEndpoint: strings.TrimSpace(azureEndpoint), 85 89 AzureOpenAIModel: strings.TrimSpace(azureDeployment), ··· 94 98 CloudflareAPIBase: strings.TrimSpace(cfg.CloudflareAPIBase), 95 99 GeminiAPIKey: strings.TrimSpace(geminiKey), 96 100 GeminiAPIBase: strings.TrimSpace(geminiBase), 101 + Pricing: pricing, 97 102 98 103 Debug: cfg.Debug, 99 104 } ··· 101 106 return &Client{ 102 107 provider: provider, 103 108 model: strings.TrimSpace(cfg.Model), 109 + pricing: pricing, 104 110 requestTimeout: cfg.RequestTimeout, 105 111 temperature: cloneFloat64(cfg.Temperature), 106 112 reasoningEffort: strings.ToLower(strings.TrimSpace(cfg.ReasoningEffort)), ··· 141 147 142 148 toolCalls := toLLMToolCalls(resp.ToolCalls) 143 149 model := firstNonEmpty(req.Model, c.model) 150 + usage := toLLMUsage(resp.Usage) 151 + if enriched, changed := enrichUsageFromOpenAICompatibleRaw(usage, resp.Raw); changed { 152 + usage = recalculateUsageCost(enriched, c.pricing, req.InferenceProvider, model) 153 + } 144 154 if shouldEnsureGeminiThoughtSignature(c.provider, model) { 145 155 toolCalls = ensureGeminiToolCallThoughtSignatures(toolCalls) 146 156 } ··· 149 159 Text: resp.Text, 150 160 Parts: toLLMParts(resp.Parts), 151 161 ToolCalls: toolCalls, 152 - Usage: toLLMUsage(resp.Usage), 162 + Usage: usage, 153 163 Duration: time.Since(start), 154 164 }, nil 155 165 } ··· 362 372 } 363 373 } 364 374 375 + type rawJSONProvider interface { 376 + RawJSON() string 377 + } 378 + 379 + func enrichUsageFromOpenAICompatibleRaw(usage llm.Usage, raw any) (llm.Usage, bool) { 380 + rawJSON := rawJSONFromOpenAICompatibleRaw(raw) 381 + if strings.TrimSpace(rawJSON) == "" { 382 + return usage, false 383 + } 384 + 385 + var payload struct { 386 + Usage struct { 387 + CachedTokens *int `json:"cached_tokens"` 388 + CacheReadInputTokens *int `json:"cache_read_input_tokens"` 389 + CacheCreationInputTokens *int `json:"cache_creation_input_tokens"` 390 + CacheCreation map[string]int `json:"cache_creation"` 391 + PromptTokensDetails struct { 392 + CachedTokens *int `json:"cached_tokens"` 393 + CacheReadInputTokens *int `json:"cache_read_input_tokens"` 394 + CacheCreationInputTokens *int `json:"cache_creation_input_tokens"` 395 + CacheCreation map[string]int `json:"cache_creation"` 396 + } `json:"prompt_tokens_details"` 397 + } `json:"usage"` 398 + } 399 + if err := json.Unmarshal([]byte(rawJSON), &payload); err != nil { 400 + return usage, false 401 + } 402 + 403 + changed := false 404 + if cached := firstPositiveInt( 405 + payload.Usage.PromptTokensDetails.CacheReadInputTokens, 406 + payload.Usage.PromptTokensDetails.CachedTokens, 407 + payload.Usage.CacheReadInputTokens, 408 + payload.Usage.CachedTokens, 409 + ); cached > 0 && usage.Cache.CachedInputTokens != cached { 410 + usage.Cache.CachedInputTokens = cached 411 + changed = true 412 + } 413 + if created := firstPositiveInt( 414 + payload.Usage.PromptTokensDetails.CacheCreationInputTokens, 415 + payload.Usage.CacheCreationInputTokens, 416 + ); created > 0 && usage.Cache.CacheCreationInputTokens != created { 417 + usage.Cache.CacheCreationInputTokens = created 418 + changed = true 419 + } 420 + var detailChanged bool 421 + usage.Cache.Details, detailChanged = mergePositiveCacheDetails(usage.Cache.Details, payload.Usage.PromptTokensDetails.CacheCreation) 422 + changed = changed || detailChanged 423 + usage.Cache.Details, detailChanged = mergePositiveCacheDetails(usage.Cache.Details, payload.Usage.CacheCreation) 424 + changed = changed || detailChanged 425 + return usage, changed 426 + } 427 + 428 + func rawJSONFromOpenAICompatibleRaw(raw any) string { 429 + if raw == nil { 430 + return "" 431 + } 432 + if v, ok := raw.(rawJSONProvider); ok { 433 + return strings.TrimSpace(v.RawJSON()) 434 + } 435 + b, err := json.Marshal(raw) 436 + if err != nil { 437 + return "" 438 + } 439 + return strings.TrimSpace(string(b)) 440 + } 441 + 442 + func firstPositiveInt(values ...*int) int { 443 + for _, value := range values { 444 + if value != nil && *value > 0 { 445 + return *value 446 + } 447 + } 448 + return 0 449 + } 450 + 451 + func mergePositiveCacheDetails(dst map[string]int, src map[string]int) (map[string]int, bool) { 452 + if len(src) == 0 { 453 + return dst, false 454 + } 455 + changed := false 456 + for key, value := range src { 457 + key = strings.TrimSpace(key) 458 + if key == "" || value <= 0 { 459 + continue 460 + } 461 + if dst == nil { 462 + dst = map[string]int{} 463 + } 464 + if dst[key] != value { 465 + dst[key] = value 466 + changed = true 467 + } 468 + } 469 + return dst, changed 470 + } 471 + 472 + func recalculateUsageCost(usage llm.Usage, pricing *uniaiapi.PricingCatalog, inferenceProvider, model string) llm.Usage { 473 + if pricing == nil { 474 + usage.Cost = nil 475 + return usage 476 + } 477 + cost, ok := pricing.EstimateChatCostWithInferenceProvider(strings.TrimSpace(inferenceProvider), strings.TrimSpace(model), uniaiapi.Usage{ 478 + InputTokens: usage.InputTokens, 479 + OutputTokens: usage.OutputTokens, 480 + TotalTokens: usage.TotalTokens, 481 + Cache: uniaiapi.UsageCache{ 482 + CachedInputTokens: usage.Cache.CachedInputTokens, 483 + CacheCreationInputTokens: usage.Cache.CacheCreationInputTokens, 484 + Details: cloneIntMap(usage.Cache.Details), 485 + }, 486 + }) 487 + if !ok { 488 + usage.Cost = nil 489 + return usage 490 + } 491 + usage.Cost = toLLMUsageCost(cost) 492 + return usage 493 + } 494 + 365 495 func cloneIntMap(in map[string]int) map[string]int { 366 496 if len(in) == 0 { 367 497 return nil ··· 624 754 } 625 755 switch strings.ToLower(rawTTL) { 626 756 case "short": 627 - return "in-memory" 757 + return "in_memory" 628 758 case "long": 629 759 return "24h" 630 760 } ··· 633 763 return "" 634 764 } 635 765 if d <= 5*time.Minute { 636 - return "in-memory" 766 + return "in_memory" 637 767 } 638 768 return "24h" 639 769 }
+77 -2
providers/uniai/client_test.go
··· 1 1 package uniai 2 2 3 3 import ( 4 + "math" 4 5 "reflect" 5 6 "testing" 6 7 ··· 228 229 } 229 230 } 230 231 232 + type testRawJSON string 233 + 234 + func (r testRawJSON) RawJSON() string { 235 + return string(r) 236 + } 237 + 238 + func TestEnrichUsageFromOpenAICompatibleRawReadsCacheCreation(t *testing.T) { 239 + usage := llm.Usage{ 240 + InputTokens: 2648, 241 + OutputTokens: 38, 242 + TotalTokens: 2686, 243 + } 244 + raw := testRawJSON(`{ 245 + "usage": { 246 + "prompt_tokens": 2648, 247 + "completion_tokens": 38, 248 + "total_tokens": 2686, 249 + "prompt_tokens_details": { 250 + "cached_tokens": 2390, 251 + "cache_read_input_tokens": 2390, 252 + "cache_creation_input_tokens": 255 253 + } 254 + } 255 + }`) 256 + 257 + got, changed := enrichUsageFromOpenAICompatibleRaw(usage, raw) 258 + if !changed { 259 + t.Fatalf("changed = false, want true") 260 + } 261 + if got.Cache.CachedInputTokens != 2390 { 262 + t.Fatalf("cached_input_tokens = %d, want 2390", got.Cache.CachedInputTokens) 263 + } 264 + if got.Cache.CacheCreationInputTokens != 255 { 265 + t.Fatalf("cache_creation_input_tokens = %d, want 255", got.Cache.CacheCreationInputTokens) 266 + } 267 + } 268 + 269 + func TestRecalculateUsageCostIncludesCacheCreation(t *testing.T) { 270 + cachedInputRate := 0.30 271 + cacheCreationRate := 3.75 272 + pricing := &uniaiapi.PricingCatalog{Chat: []uniaiapi.ChatPricingRule{ 273 + { 274 + Model: "claude-sonnet-4-6", 275 + InputUSDPerMillion: 3.0, 276 + OutputUSDPerMillion: 15.0, 277 + CachedInputUSDPerMillion: &cachedInputRate, 278 + CacheCreationInputUSDPerMillion: &cacheCreationRate, 279 + CacheCreationInputDetailUSDPerMillion: nil, 280 + }, 281 + }} 282 + usage := llm.Usage{ 283 + InputTokens: 2648, 284 + OutputTokens: 38, 285 + TotalTokens: 2686, 286 + Cache: llm.UsageCache{ 287 + CachedInputTokens: 2390, 288 + CacheCreationInputTokens: 255, 289 + }, 290 + Cost: &llm.UsageCost{Total: 999}, 291 + } 292 + 293 + got := recalculateUsageCost(usage, pricing, "", "anthropic/claude-sonnet-4-6") 294 + if got.Cost == nil { 295 + t.Fatalf("cost = nil") 296 + } 297 + wantTotal := 0.00225225 298 + if math.Abs(got.Cost.Total-wantTotal) > 0.000000001 { 299 + t.Fatalf("total cost = %.10f, want %.10f", got.Cost.Total, wantTotal) 300 + } 301 + if got.Cost.CacheCreationInput <= 0 { 302 + t.Fatalf("cache creation cost = %.10f, want > 0", got.Cost.CacheCreationInput) 303 + } 304 + } 305 + 231 306 func TestBuildChatOptionsDisablesOnStreamForGeminiProvider(t *testing.T) { 232 307 req := llm.Request{ 233 308 Messages: []llm.Message{{Role: "user", Content: "hello"}}, ··· 363 438 if got := built.Options.OpenAI["prompt_cache_key"]; got == "" || got == nil { 364 439 t.Fatalf("prompt_cache_key = %#v, want non-empty derived key", got) 365 440 } 366 - if got := built.Options.OpenAI["prompt_cache_retention"]; got != "in-memory" { 367 - t.Fatalf("prompt_cache_retention = %#v, want in-memory", got) 441 + if got := built.Options.OpenAI["prompt_cache_retention"]; got != "in_memory" { 442 + t.Fatalf("prompt_cache_retention = %#v, want in_memory", got) 368 443 } 369 444 if got := built.Options.OpenAI["response_format"]; got != "json_object" { 370 445 t.Fatalf("response_format = %#v, want json_object", got)