diff --git a/internal/router/budget_fallback.go b/internal/router/budget_fallback.go index bc3cb9d3..66c8aefe 100644 --- a/internal/router/budget_fallback.go +++ b/internal/router/budget_fallback.go @@ -41,11 +41,11 @@ func (r *Router) budgetFallback(p auth.Principal, req *parsedRequest) (accountin // messages to stay under a limit they did not ask about is a worse // answer than saying the prompt is too large, and it would be done // silently, to a request that was already being quietly moved. - _, cand := r.resolve(name) + canonical, cand := r.resolve(name) if why := checkContextFits(cand, req.PromptTokens, req.MaxTokens); why != "" { continue } - res, err := r.acct.Check(p, name, req.estimatedTokens()) + res, err := r.acct.Check(p, canonical, req.estimatedTokens()) if err != nil { continue } diff --git a/internal/router/handlers.go b/internal/router/handlers.go index 267e1ecd..6c2f52ae 100644 --- a/internal/router/handlers.go +++ b/internal/router/handlers.go @@ -282,7 +282,13 @@ func (r *Router) handleInference(path string) http.HandlerFunc { // over its quota should be refused rather than queued, and queueing // it would have it occupy a band's capacity to no purpose. if r.acct != nil { - reservation, err := r.acct.Check(p, parsed.Model, parsed.estimatedTokens()) + // Priced by the model that will serve, not by the name the + // caller typed. A price table is written against the catalog, + // and an alias asked for by its short name would miss every + // entry and fall to the default — charging a 1.5B at whatever + // rate the largest model set, or at nothing. + priced, _ := r.resolve(parsed.Model) + reservation, err := r.acct.Check(p, priced, parsed.estimatedTokens()) if err != nil && isBudgetSpent(err) { // Out of money, not out of options. A smaller model the // caller can still afford beats a refusal, and the catalog diff --git a/internal/router/pricing_test.go b/internal/router/pricing_test.go new file mode 100644 index 00000000..9135cf87 --- /dev/null +++ b/internal/router/pricing_test.go @@ -0,0 +1,142 @@ +package router + +import ( + "crypto/sha256" + "path/filepath" + "strings" + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// pricedRouter builds a gateway with authentication, one model under an alias, +// and a price table that charges that model and almost nothing for anything +// else — so a request priced by the wrong name is visible in the figure rather +// than merely plausible. +func pricedRouter(t *testing.T, engineURL string, served ...string) (*Router, string) { + t.Helper() + + cfg := baseConfig(config.Endpoint{ + ID: "a", URL: engineURL, + Models: append([]string{"qwen2.5-14b-instruct"}, served...), + }) + cfg.Models = []catalog.Model{{ + Name: "qwen2.5-14b-instruct", Aliases: []string{"q14"}, + }} + cfg.Auth.Enabled = true + cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db") + cfg.Auth.Pricing = map[string]config.Price{ + "qwen2.5-14b-instruct": {PromptPer1k: 14000, CompletionPer1k: 42000}, + } + // Deliberately not zero: a fallback of nothing would make a miss look the + // same as a model that is free. + cfg.Auth.DefaultPrice = config.Price{PromptPer1k: 1, CompletionPer1k: 1} + cfg.Auth.FlushInterval = 20 * time.Millisecond + cfg.Auth.BatchSize = 1 + + rt := newTestRouter(t, cfg) + + org, err := rt.store.CreateOrg(t.Context(), "acme") + if err != nil { + t.Fatal(err) + } + team, err := rt.store.CreateTeam(t.Context(), org, "platform") + if err != nil { + t.Fatal(err) + } + project, err := rt.store.CreateProject(t.Context(), team, "gateway") + if err != nil { + t.Fatal(err) + } + + // A real key shape: authentication rejects anything else before it looks. + secret := store.KeyPrefix + strings.Repeat("a", 43) + sum := sha256.Sum256([]byte(secret)) + if err := rt.store.CreateKey(t.Context(), store.Key{ + ID: "key-1", Hash: sum[:], Prefix: "aaaa", Name: "test", + ProjectID: project, CreatedAt: time.Now(), + }); err != nil { + t.Fatal(err) + } + return rt, secret +} + +func TestAnAliasIsPricedAsTheModelItResolvesTo(t *testing.T) { + // The price table is written against the catalog, and a caller asks for + // whichever name they were given. Charging by the string they typed meant + // every alias missed the table and fell to the default — so a budget could + // be spent at the wrong rate, or not at all, depending on which name went + // into the request. + e := newFakeEngine(t, "qwen2.5-14b-instruct") + rt, secret := pricedRouter(t, e.srv.URL) + + rec := post(t, rt.Handler(), + map[string]any{"model": "q14", "messages": []any{ + map[string]any{"role": "user", "content": "hello"}}}, + map[string]string{"Authorization": "Bearer " + secret}) + if rec.Code != 200 { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + + u := waitForUsage(t, rt) + want := charge(u.Prompt, 14000) + charge(u.Completion, 42000) + fallback := charge(u.Prompt, 1) + charge(u.Completion, 1) + if u.CostMicros != want { + t.Errorf("charged %d micros for %d+%d tokens, expected %d; the "+ + "fallback rate would have charged %d", + u.CostMicros, u.Prompt, u.Completion, want, fallback) + } +} + +// charge is Price.Cost for one half of a request, rounding up the way the +// accountant does — spelled out here so the expected figure is arithmetic the +// test states rather than the code under test restated. +func charge(tokens, per1k int64) int64 { + if tokens <= 0 { + return 0 + } + return (tokens*per1k + 999) / 1000 +} + +func TestAModelWithNoPriceFallsBackRatherThanCostingNothing(t *testing.T) { + e := newFakeEngine(t, "unpriced") + rt, secret := pricedRouter(t, e.srv.URL, "unpriced") + + // Not in the catalog and not in the price table. It must still cost the + // fallback, because a model nobody priced is the one most likely to be new. + rec := post(t, rt.Handler(), + map[string]any{"model": "unpriced", "messages": []any{ + map[string]any{"role": "user", "content": "hello"}}}, + map[string]string{"Authorization": "Bearer " + secret}) + if rec.Code != 200 { + t.Fatalf("status %d: %s", rec.Code, rec.Body.String()) + } + + u := waitForUsage(t, rt) + want := charge(u.Prompt, 1) + charge(u.Completion, 1) + if u.CostMicros != want { + t.Errorf("charged %d micros for %d+%d tokens, expected the fallback's %d", + u.CostMicros, u.Prompt, u.Completion, want) + } + if u.CostMicros == 0 { + t.Error("an unpriced model cost nothing; the fallback did not apply") + } +} + +func waitForUsage(t *testing.T, rt *Router) store.KeyUsage { + t.Helper() + since := time.Now().Add(-time.Hour) + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + rows, err := rt.store.UsageSince(t.Context(), since) + if err == nil && len(rows) > 0 { + return rows[0] + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("no usage was recorded") + return store.KeyUsage{} +}