hobokenchicken
90874a6721
Removed .env and .env.backup from git tracking and consolidated configuration into .env.example. Updated .gitignore to robustly prevent accidental inclusion of sensitive files.
55 lines
1.4 KiB
Go
55 lines
1.4 KiB
Go
package utils
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"time"
|
|
|
|
"llm-proxy/internal/models"
|
|
"github.com/go-resty/resty/v2"
|
|
)
|
|
|
|
const ModelsDevURL = "https://models.dev/api.json"
|
|
|
|
func FetchRegistry() (*models.ModelRegistry, error) {
|
|
log.Printf("Fetching model registry from %s", ModelsDevURL)
|
|
|
|
client := resty.New().SetTimeout(10 * time.Second)
|
|
resp, err := client.R().Get(ModelsDevURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to fetch registry: %w", err)
|
|
}
|
|
|
|
if !resp.IsSuccess() {
|
|
return nil, fmt.Errorf("failed to fetch registry: HTTP %d", resp.StatusCode())
|
|
}
|
|
|
|
var providers map[string]models.ProviderInfo
|
|
if err := json.Unmarshal(resp.Body(), &providers); err != nil {
|
|
return nil, fmt.Errorf("failed to unmarshal registry: %w", err)
|
|
}
|
|
|
|
log.Println("Successfully loaded model registry")
|
|
return &models.ModelRegistry{Providers: providers}, nil
|
|
}
|
|
|
|
func CalculateCost(registry *models.ModelRegistry, modelID string, promptTokens, completionTokens, cacheRead, cacheWrite uint32) float64 {
|
|
meta := registry.FindModel(modelID)
|
|
if meta == nil || meta.Cost == nil {
|
|
return 0.0
|
|
}
|
|
|
|
cost := (float64(promptTokens) * meta.Cost.Input / 1000000.0) +
|
|
(float64(completionTokens) * meta.Cost.Output / 1000000.0)
|
|
|
|
if meta.Cost.CacheRead != nil {
|
|
cost += float64(cacheRead) * (*meta.Cost.CacheRead) / 1000000.0
|
|
}
|
|
if meta.Cost.CacheWrite != nil {
|
|
cost += float64(cacheWrite) * (*meta.Cost.CacheWrite) / 1000000.0
|
|
}
|
|
|
|
return cost
|
|
}
|