From dfb949d4372eba86fab382a5038c56b86a0582b5 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:24:47 +0000 Subject: [PATCH 01/82] fix(terraform)!: send additional_drop_params to the proxy instead of deleting the keys it names (#46782) * fix(terraform): send additional_drop_params to the proxy instead of deleting the keys it names litellm_model treated additional_drop_params in additional_litellm_params as a provider-side delete list: it never sent the list and removed every key it named from litellm_params, so vertex_location vanished and requests still carried store. The list now goes to the proxy like any other JSON value, so the proxy drops those params per request. Approach credited to #38345. Fixes #38349 * fix(terraform)!: validate additional_drop_params and read it back from the proxy on refresh * docs(terraform): note that a non-list additional_drop_params now fails the apply * fix(terraform)!: check additional_drop_params at plan time and refresh listed keys and outside lists * fix(terraform)!: show every listed key an older provider deleted in the upgrade plan and leave outside drop lists alone Refresh no longer reads listed keys back from the proxy, which left a plan diff that never cleared when a listed key was removed from config. When the proxy has no stored list, the plan now shows the list and each listed key the config sets as changes. A drop list set outside Terraform on a model whose config has none stays out of the plan, and an empty value now fails the plan * fix(terraform): show a listed merge_reasoning_content_in_choices in the upgrade plan The refresh now clears merge_reasoning_content_in_choices when it is listed and the proxy has no stored list, since an older provider deleted it, and no longer clears litellm_credential_name, which the older provider always sent --- terraform/provider/CHANGELOG.md | 2 + terraform/provider/docs/resources/model.md | 14 +- .../examples/model_additional_params.tf | 21 +- terraform/provider/litellm/resource_model.go | 12 +- .../provider/litellm/resource_model_crud.go | 192 +++++++++----- .../provider/litellm/resource_model_test.go | 243 ++++++++++++++++++ terraform/provider/litellm/types.go | 1 + 7 files changed, 394 insertions(+), 91 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 4282485714a..597e0c6d4d8 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -51,6 +51,7 @@ longer signal it. ### Fixed +- **model**: `additional_drop_params` in `litellm_model.additional_litellm_params` is now sent to the proxy as the deployment's `additional_drop_params` list, the same list a `model_list` entry takes. Before, the provider never sent the list and instead deleted every key it named from `litellm_params`, so a `store` or `vertex_location` listed there was lost and requests still carried the parameters (#38349). The value must be a JSON list of parameter names such as `jsonencode(["store"])`, which `terraform plan` now checks, and `null` entries are skipped as before. When the configuration sets the list, refresh reads the stored list back from the proxy, so a list changed outside Terraform shows up as drift. On a model an older provider created, the proxy has no stored list, so the first plan after upgrading shows the list and every listed key the configuration sets as changes, since that apply writes them. A write-only secret such as `model_api_key_wo` is never shown in a plan - **user**: `litellm_user` updates no longer send `blocked`, which `/user/update` writes into the user table where no such column exists, so every update failed with a 400 (`Could not find field at upsertOneLiteLLM_UserTable.create.blocked`). Create still sends `blocked`, including `false`, so a `default_internal_user_params.blocked` on the proxy cannot block the auto-created key. Changing `blocked` after creation has no effect, since the proxy has no user-level block; use `litellm_key_block` instead - **model**: `litellm_model` refresh now reads the `{"data": [...]}` envelope `/model/info` returns, so `model_info` fields changed outside Terraform show up as drift instead of silently keeping the previous state - **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update @@ -69,6 +70,7 @@ longer signal it. ### Changed +- **model** (breaking): a key named in `additional_litellm_params.additional_drop_params` is no longer deleted from the deployment's `litellm_params`. A configuration that also sets that key, such as `vertex_location`, stores it on the deployment at its next apply, and the proxy applies the drop list to requests instead. For `vertex_location` this changes routing. `/v1/messages` calls now go to the configured region, while chat completions drop a listed `vertex_location` and use the default region. To keep such a key off a deployment, remove its attribute and replace the model with `terraform apply -replace`, which gives it a new model id, since an update keeps a stored value the provider leaves out. `model`, `custom_llm_provider` and `merge_reasoning_content_in_choices` are always sent. A value that is not a JSON list of names, such as a bare `"store"` or an empty string, now fails the plan, and with it `terraform destroy`, instead of being stored as a string the proxy ignores. Rewrite it as `jsonencode([...])`, or remove the model without fixing it through `terraform destroy -refresh=false` - **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index b7c440b0e65..85c495bf1ca 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -183,16 +183,16 @@ The following arguments are supported: * Numeric strings are parsed first as integers; if integer parsing fails, parsed as floats (e.g., `"16384"` -> 16384, `"0.75"` -> 0.75) * JSON strings (starting with `[` or `{`) are parsed as JSON objects/arrays * Non-convertible strings remain strings - * Non-string map values (if supplied) are passed through unchanged. * The provider merges these keys into the `litellm_params` payload sent to the API. * Note: the remote API may not echo back all custom parameters; this provider preserves `additional_litellm_params` in state when present in configuration. - **Special parameter: `additional_drop_params`** - * When `additional_drop_params` is provided as a JSON array string, it specifies parameters to remove from the final `litellm_params` before sending to the API - * This allows you to override or remove built-in parameters if needed - * The `additional_drop_params` key itself is not included in the final parameters + **`additional_drop_params`** + * Give it a JSON array string, such as `jsonencode(["store"])`. The provider sends the list to the proxy as the deployment's `additional_drop_params`, and the proxy applies it the same as `additional_drop_params` in a `model_list` entry. `terraform plan` fails on any other value, an empty value, or an entry that is not a name, and `null` entries are skipped + * Other keys on the resource, such as `vertex_location`, are sent unchanged + * When the configuration sets the list, refresh reads the stored list back from the proxy, so a list changed outside Terraform shows up as drift. A model whose configuration has no list keeps a list set outside Terraform out of the plan, but the proxy clears it on the next update + * When the proxy has no stored list, as on a model an older provider created, the next plan shows the list and every listed key the configuration sets as changes, since the next apply writes them - Example showing booleans, integers, floats, strings, and parameter dropping: + Example showing booleans, integers, floats, strings, JSON, and dropped request parameters: ```hcl resource "litellm_model" "with_additional" { @@ -208,7 +208,7 @@ The following arguments are supported: "scale" = "0.75" # becomes float 0.75 "note" = "for testing" # stays string "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object - "additional_drop_params" = "[\"reasoningEffort\"]" # removes reasoningEffort parameter + "additional_drop_params" = "[\"reasoning_effort\"]" # proxy drops reasoning_effort from requests } } ``` diff --git a/terraform/provider/examples/model_additional_params.tf b/terraform/provider/examples/model_additional_params.tf index fb4981ec6fb..64c37993e04 100644 --- a/terraform/provider/examples/model_additional_params.tf +++ b/terraform/provider/examples/model_additional_params.tf @@ -9,8 +9,7 @@ provider "litellm" { # - Numeric strings will be parsed to integer (if possible) otherwise float. # - JSON strings (starting with [ or {) will be parsed as JSON objects/arrays. # - Non-convertible strings remain strings. -# - Non-string map values are passed through unchanged. -# - Use "additional_drop_params" as a JSON array to remove parameters from the final request. +# - "additional_drop_params" takes a JSON array of request parameters the proxy drops before calling the model provider. resource "litellm_model" "with_additional" { model_name = "custom-model" @@ -21,14 +20,12 @@ resource "litellm_model" "with_additional" { # Additional parameters not exposed as first-class arguments additional_litellm_params = { - "use_fine_tune" = "true" # becomes boolean true - "max_context" = "16384" # becomes integer 16384 - "temperature_scale" = "0.75" # becomes float 0.75 - "experimental_feature" = "enabled" # stays string "enabled" - "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object - "additional_drop_params" = "[\"reasoningEffort\"]" # removes reasoningEffort parameter - # You may also pass non-string values (they will be passed through unchanged) - # "raw_flag" = true + "use_fine_tune" = "true" # becomes boolean true + "max_context" = "16384" # becomes integer 16384 + "temperature_scale" = "0.75" # becomes float 0.75 + "experimental_feature" = "enabled" # stays string "enabled" + "complex_config" = "{\"nested\": {\"value\": 42}}" # parsed as JSON object + "additional_drop_params" = "[\"reasoning_effort\"]" # proxy drops reasoning_effort from requests } # Cost configuration (optional) @@ -47,9 +44,9 @@ resource "litellm_model" "azure_with_drop_params" { tier = "paid" mode = "completion" - # Drop the reasoningEffort parameter that might be automatically added + # Have the proxy drop reasoning_effort from requests before they reach Azure additional_litellm_params = { - "additional_drop_params" = "[\"reasoningEffort\"]" + "additional_drop_params" = "[\"reasoning_effort\"]" } input_cost_per_million_tokens = 0.25 diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 1dc029b861c..e9582ad75fb 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -7,10 +7,11 @@ import ( func resourceLiteLLMModel() *schema.Resource { return &schema.Resource{ - Create: resourceLiteLLMModelCreate, - Read: resourceLiteLLMModelRead, - Update: resourceLiteLLMModelUpdate, - Delete: resourceLiteLLMModelDelete, + Create: resourceLiteLLMModelCreate, + Read: resourceLiteLLMModelRead, + Update: resourceLiteLLMModelUpdate, + Delete: resourceLiteLLMModelDelete, + CustomizeDiff: validateAdditionalDropParams, Importer: &schema.ResourceImporter{ StateContext: schema.ImportStatePassthroughContext, }, @@ -237,7 +238,8 @@ func resourceLiteLLMModel() *schema.Resource { Elem: &schema.Schema{ Type: schema.TypeString, }, - Description: "Additional parameters to pass to litellm_params beyond the standard ones", + DiffSuppressFunc: suppressEquivalentDropParams, + Description: "Additional parameters to pass to litellm_params beyond the standard ones", }, }, } diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 3309dbe26dd..3db6f68760d 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -1,10 +1,13 @@ package litellm import ( + "context" "encoding/json" "fmt" "log" + "maps" "net/url" + "slices" "strconv" "strings" "time" @@ -68,6 +71,119 @@ func secretFromResource(d *schema.ResourceData, attr string) string { return d.Get(attr).(string) } +const additionalDropParamsKey = "additional_drop_params" + +func parseAdditionalDropParams(raw string) ([]string, error) { + var items []interface{} + if err := json.Unmarshal([]byte(raw), &items); err != nil { + return nil, fmt.Errorf("additional_litellm_params.%s must be a JSON list of parameter names, such as jsonencode([\"store\"]): %w", additionalDropParamsKey, err) + } + names := []string{} + for _, item := range items { + if item == nil { + continue + } + name, ok := item.(string) + if !ok { + return nil, fmt.Errorf("additional_litellm_params.%s entries must be parameter names, got %v", additionalDropParamsKey, item) + } + names = append(names, name) + } + return names, nil +} + +func validateAdditionalDropParams(_ context.Context, d *schema.ResourceDiff, _ interface{}) error { + params, _ := d.Get("additional_litellm_params").(map[string]interface{}) + raw, set := params[additionalDropParamsKey].(string) + if !set { + return nil + } + _, err := parseAdditionalDropParams(raw) + return err +} + +func suppressEquivalentDropParams(key, oldValue, newValue string, _ *schema.ResourceData) bool { + if key != "additional_litellm_params."+additionalDropParamsKey { + return false + } + oldNames, oldErr := parseAdditionalDropParams(oldValue) + newNames, newErr := parseAdditionalDropParams(newValue) + return oldErr == nil && newErr == nil && slices.Equal(oldNames, newNames) +} + +var attributeSentAsLiteLLMParam = map[string]string{ + "custom_llm_provider": "custom_llm_provider", + "tpm": "tpm", + "rpm": "rpm", + "input_cost_per_token": "input_cost_per_million_tokens", + "output_cost_per_token": "output_cost_per_million_tokens", + "api_key": "model_api_key", + "api_base": "model_api_base", + "api_version": "api_version", + "input_cost_per_pixel": "input_cost_per_pixel", + "output_cost_per_pixel": "output_cost_per_pixel", + "input_cost_per_second": "input_cost_per_second", + "output_cost_per_second": "output_cost_per_second", + "aws_access_key_id": "aws_access_key_id", + "aws_secret_access_key": "aws_secret_access_key", + "aws_region_name": "aws_region_name", + "aws_session_name": "aws_session_name", + "aws_role_name": "aws_role_name", + "vertex_project": "vertex_project", + "vertex_location": "vertex_location", + "vertex_credentials": "vertex_credentials", + "reasoning_effort": "reasoning_effort", + "thinking": "thinking_enabled", + "merge_reasoning_content_in_choices": "merge_reasoning_content_in_choices", +} + +func refreshAdditionalDropParams(d *schema.ResourceData, stored interface{}) { + params, _ := d.Get("additional_litellm_params").(map[string]interface{}) + raw, managed := params[additionalDropParamsKey].(string) + if !managed { + return + } + refreshed := maps.Clone(params) + if stored != nil { + encoded, _ := json.Marshal(stored) + refreshed[additionalDropParamsKey] = string(encoded) + d.Set("additional_litellm_params", refreshed) + return + } + delete(refreshed, additionalDropParamsKey) + listed, _ := parseAdditionalDropParams(raw) + for _, key := range listed { + delete(refreshed, key) + if attribute, ok := attributeSentAsLiteLLMParam[key]; ok { + d.Set(attribute, nil) + } + } + d.Set("additional_litellm_params", refreshed) +} + +func coerceAdditionalParamValue(strValue string) interface{} { + trimmedValue := strings.TrimSpace(strValue) + if strings.HasPrefix(trimmedValue, "[") || strings.HasPrefix(trimmedValue, "{") { + var parsedValue interface{} + if err := json.Unmarshal([]byte(strValue), &parsedValue); err == nil { + return parsedValue + } + } + switch strValue { + case "true": + return true + case "false": + return false + } + if intValue, err := strconv.Atoi(strValue); err == nil { + return intValue + } + if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { + return floatValue + } + return strValue +} + func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) error { client, ok := m.(*Client) if !ok { @@ -175,76 +291,17 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e litellmParams["thinking"] = thinking } - // Add additional parameters if provided if additionalParams, ok := d.GetOk("additional_litellm_params"); ok { - var dropParams []string - for key, value := range additionalParams.(map[string]interface{}) { - // Convert string values to appropriate types where possible - if strValue, ok := value.(string); ok { - // Check if it's JSON (starts with [ or {) - trimmedValue := strings.TrimSpace(strValue) - if strings.HasPrefix(trimmedValue, "[") || strings.HasPrefix(trimmedValue, "{") { - var parsedValue interface{} - if err := json.Unmarshal([]byte(strValue), &parsedValue); err == nil { - // Successfully parsed JSON - if key == "additional_drop_params" { - // Handle drop params specially - if dropList, ok := parsedValue.([]interface{}); ok { - for _, item := range dropList { - if paramStr, ok := item.(string); ok { - dropParams = append(dropParams, paramStr) - } - } - } - continue // Don't add to litellmParams - } else { - litellmParams[key] = parsedValue - } - } else { - // Not valid JSON, apply existing conversion logic - if strValue == "true" { - litellmParams[key] = true - } else if strValue == "false" { - litellmParams[key] = false - } else { - // Try to convert numeric strings - if intValue, err := strconv.Atoi(strValue); err == nil { - litellmParams[key] = intValue - } else if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { - litellmParams[key] = floatValue - } else { - // Keep as string - litellmParams[key] = strValue - } - } - } - } else { - // Apply existing conversion logic for non-JSON strings - if strValue == "true" { - litellmParams[key] = true - } else if strValue == "false" { - litellmParams[key] = false - } else { - // Try to convert numeric strings - if intValue, err := strconv.Atoi(strValue); err == nil { - litellmParams[key] = intValue - } else if floatValue, err := strconv.ParseFloat(strValue, 64); err == nil { - litellmParams[key] = floatValue - } else { - // Keep as string - litellmParams[key] = strValue - } - } - } - } else { - litellmParams[key] = value + if key != additionalDropParamsKey { + litellmParams[key] = coerceAdditionalParamValue(value.(string)) + continue } - } - - // Apply drop params at the end - for _, paramToDrop := range dropParams { - delete(litellmParams, paramToDrop) + dropParams, err := parseAdditionalDropParams(value.(string)) + if err != nil { + return err + } + litellmParams[key] = dropParams } } @@ -419,6 +476,7 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { if _, ok := d.GetOk("additional_litellm_params"); ok { d.Set("additional_litellm_params", d.Get("additional_litellm_params")) } + refreshAdditionalDropParams(d, modelResp.LiteLLMParams.AdditionalDropParams) return nil } diff --git a/terraform/provider/litellm/resource_model_test.go b/terraform/provider/litellm/resource_model_test.go index 22c226adfd3..4c2d5d27442 100644 --- a/terraform/provider/litellm/resource_model_test.go +++ b/terraform/provider/litellm/resource_model_test.go @@ -3,8 +3,11 @@ package litellm import ( "context" "encoding/json" + "maps" "net/http" "net/http/httptest" + "reflect" + "slices" "testing" "github.com/hashicorp/go-cty/cty" @@ -116,6 +119,246 @@ func TestResourceLiteLLMModelCreateOmitsUnsetDisplayName(t *testing.T) { } } +func TestResourceLiteLLMModelCreateSendsAdditionalLiteLLMParams(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "vertex_project": "test-project", + "vertex_location": "eu", + "additional_litellm_params": map[string]interface{}{ + "additional_drop_params": `["store", "vertex_location"]`, + "tags": ` ["prod", "eu"]`, + "headers": `{"x-team": "search"}`, + "use_in_pass_through": "true", + "merge_reasoning_content_in_choices": "false", + "max_parallel_requests": "8", + "input_cost_per_token": "0.5", + "description": "[draft] haiku", + }, + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + litellmParams, ok := createPayload["litellm_params"].(map[string]interface{}) + if !ok { + t.Fatalf("expected litellm_params object in create payload, got %v", createPayload["litellm_params"]) + } + want := map[string]interface{}{ + "additional_drop_params": []interface{}{"store", "vertex_location"}, + "vertex_project": "test-project", + "vertex_location": "eu", + "tags": []interface{}{"prod", "eu"}, + "headers": map[string]interface{}{"x-team": "search"}, + "use_in_pass_through": true, + "merge_reasoning_content_in_choices": false, + "max_parallel_requests": float64(8), + "input_cost_per_token": 0.5, + "description": "[draft] haiku", + } + for key, wantValue := range want { + if got := litellmParams[key]; !reflect.DeepEqual(got, wantValue) { + t.Errorf("litellm_params[%q] = %#v, want %#v", key, got, wantValue) + } + } +} + +const terraformUnknownValue = "74D93920-ED26-11E3-AC10-0800200C9A66" + +func TestResourceLiteLLMModelChecksAdditionalDropParams(t *testing.T) { + cases := map[string]struct { + dropParams string + want []interface{} + wantErr bool + }{ + "null entry from an unset variable is skipped": {dropParams: `["store", null]`, want: []interface{}{"store"}}, + "number entry is rejected": {dropParams: `["store", 1]`, wantErr: true}, + "object entry is rejected": {dropParams: `["store", {"name": "temperature"}]`, wantErr: true}, + "object is rejected": {dropParams: `{"store": true}`, wantErr: true}, + "bare name is rejected": {dropParams: "store", wantErr: true}, + "empty value is rejected": {dropParams: "", wantErr: true}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + res := resourceLiteLLMModel() + config := map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "additional_litellm_params": map[string]interface{}{"additional_drop_params": tc.dropParams}, + } + if _, err := res.Diff(context.Background(), nil, terraform.NewResourceConfigRaw(config), nil); (err != nil) != tc.wantErr { + t.Errorf("plan error = %v, want error %v", err, tc.wantErr) + } + + d := schema.TestResourceDataRaw(t, res.Schema, config) + err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)) + if tc.wantErr { + if err == nil { + t.Fatal("expected create to fail") + } + if createPayload != nil { + t.Errorf("expected no /model/new request, got %v", createPayload) + } + return + } + if err != nil { + t.Fatalf("create failed: %v", err) + } + litellmParams := createPayload["litellm_params"].(map[string]interface{}) + if got := litellmParams["additional_drop_params"]; !reflect.DeepEqual(got, tc.want) { + t.Errorf("additional_drop_params = %#v, want %#v", got, tc.want) + } + }) + } +} + +func TestResourceLiteLLMModelPlanAcceptsDropParamsUnknownUntilApply(t *testing.T) { + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "additional_litellm_params": map[string]interface{}{"additional_drop_params": terraformUnknownValue}, + }) + if _, err := resourceLiteLLMModel().Diff(context.Background(), nil, config, nil); err != nil { + t.Errorf("plan failed on a drop list not known until apply: %v", err) + } +} + +func TestResourceLiteLLMModelPlanComparesAdditionalDropParamsWithProxy(t *testing.T) { + const ( + listKey = "additional_litellm_params.additional_drop_params" + parallelKey = "additional_litellm_params.max_parallel_requests" + ) + cases := map[string]struct { + configured string + configLocation string + stored interface{} + storedLocation string + extraConfig map[string]interface{} + wantChanged []string + }{ + "list never sent by an older provider": {configured: `["store"]`, configLocation: "eu", storedLocation: "eu", wantChanged: []string{listKey}}, + "listed keys an older provider deleted": { + configured: `["store", "vertex_location", "tpm", "max_parallel_requests"]`, + configLocation: "eu", + wantChanged: []string{listKey, "vertex_location", "tpm", parallelKey}, + }, + "listed flag an older provider deleted": { + configured: `["merge_reasoning_content_in_choices"]`, + configLocation: "eu", + storedLocation: "eu", + extraConfig: map[string]interface{}{"merge_reasoning_content_in_choices": true}, + wantChanged: []string{listKey, "merge_reasoning_content_in_choices"}, + }, + "listed key an older provider still sent": { + configured: `["litellm_credential_name"]`, + configLocation: "eu", + storedLocation: "eu", + extraConfig: map[string]interface{}{"litellm_credential_name": "vertex-cred"}, + wantChanged: []string{listKey}, + }, + "list changed in the Admin UI": {configured: `["store"]`, configLocation: "eu", stored: []string{"temperature"}, storedLocation: "eu", wantChanged: []string{listKey}}, + "list set outside Terraform": {configLocation: "eu", stored: []string{"store"}, storedLocation: "eu"}, + "same list written with spaces": {configured: `["store", "vertex_location"]`, configLocation: "eu", stored: []string{"store", "vertex_location"}, storedLocation: "eu"}, + "null entry the proxy never received": {configured: `["store", null]`, configLocation: "eu", stored: []string{"store"}, storedLocation: "eu"}, + "listed key left out of the config": {configured: `["vertex_location"]`, stored: []string{"vertex_location"}, storedLocation: "us-east5"}, + "no list on the model or in the proxy": {configLocation: "eu", storedLocation: "eu"}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + litellmParams := map[string]interface{}{"model": "vertex_ai/claude-haiku-4-5", "custom_llm_provider": "vertex_ai"} + if tc.stored != nil { + litellmParams["additional_drop_params"] = tc.stored + } + if tc.storedLocation != "" { + litellmParams["vertex_location"] = tc.storedLocation + } + body, _ := json.Marshal(map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "litellm_params": litellmParams, + "model_info": map[string]interface{}{"id": "model-123", "db_model": true, "base_model": "claude-haiku-4-5"}, + }) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write(body) + })) + defer srv.Close() + + res := resourceLiteLLMModel() + additionalParams := map[string]interface{}{"max_parallel_requests": "8"} + if tc.configured != "" { + additionalParams["additional_drop_params"] = tc.configured + } + config := map[string]interface{}{ + "model_name": "vertex-claude-haiku", + "custom_llm_provider": "vertex_ai", + "base_model": "claude-haiku-4-5", + "tpm": 1000, + "additional_litellm_params": additionalParams, + } + if tc.configLocation != "" { + config["vertex_location"] = tc.configLocation + } + maps.Copy(config, tc.extraConfig) + d := schema.TestResourceDataRaw(t, res.Schema, config) + d.SetId("model-123") + if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + diff, err := res.Diff(context.Background(), d.State(), terraform.NewResourceConfigRaw(config), nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + for _, key := range []string{listKey, "vertex_location", "tpm", parallelKey, "merge_reasoning_content_in_choices", "litellm_credential_name"} { + attr := (*terraform.ResourceAttrDiff)(nil) + if diff != nil { + attr = diff.Attributes[key] + } + got := attr != nil && attr.Old != attr.New + if want := slices.Contains(tc.wantChanged, key); got != want { + t.Errorf("plan shows %s change = %v, want %v", key, got, want) + } + } + }) + } +} + func TestResourceLiteLLMModelReadDisplayName(t *testing.T) { cases := map[string]struct { serverBody string diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 3af4650de53..aff90cd64be 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -120,6 +120,7 @@ type LiteLLMParams struct { VertexProject string `json:"vertex_project,omitempty"` VertexLocation string `json:"vertex_location,omitempty"` VertexCredentials string `json:"vertex_credentials,omitempty"` + AdditionalDropParams interface{} `json:"additional_drop_params,omitempty"` } // ModelInfo represents information about a model. From dba86bec28cb483c9937470c5d950f5e03699d3b Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:32:24 +0000 Subject: [PATCH 02/82] fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them (#46767) * fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does * test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends * fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types * fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them * test(count_tokens): compare tools by value in the public api test * fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field. * test(count_tokens): expect the image detail the adapter now sends after merging main * fix(count_tokens): fetch Gemini remote media without blocking the event loop * fix(count_tokens): type normalize_native_messages without the unimported Any On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError * chore(count_tokens): give the Gemini count helpers' list returns a reason * fix(count_tokens): normalize Anthropic messages for Gemini counts the way /v1/messages does * fix(count_tokens): count Gemini bodies after inference's history rewrites and answer converter errors with 400 * test(count_tokens): cover the Gemini count error paths and compaction with a text-block system --------- Co-authored-by: Mrinal Chanshetty --- .../pass_through/adapters/transformation.py | 2 +- litellm/llms/gemini/common_utils.py | 26 +- litellm/llms/vertex_ai/common_utils.py | 25 +- .../llms/vertex_ai/count_tokens/context.py | 114 ++++++- .../llms/vertex_ai/gemini/transformation.py | 8 +- .../test_vertex_gemini_count_tokens_wire.py | 87 +++++ .../llms/gemini/test_gemini_common_utils.py | 98 ++++-- .../vertex_ai/count_tokens/test_context.py | 315 +++++++++++++++++- .../vertex_ai/test_vertex_ai_common_utils.py | 91 ++++- 9 files changed, 703 insertions(+), 63 deletions(-) diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 98cd16c6403..d1560c5b2ed 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -465,7 +465,7 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, - ) -> list: + ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) leading_count: Final = next( diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 57ee0f33879..bac5824b981 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -504,25 +504,41 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, vertex_gemini_count_context, ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) + resolved_contents: Final = ( + contents + if contents is not None + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request + ) + ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, "system_instruction": count_context.system_instruction, "tools": count_context.tools, } count_tokens_params_request.update(count_tokens_params) - result: Final = await GoogleAIStudioTokenCounter().acount_tokens( - **count_tokens_params_request, - ) + try: + result: Final = await GoogleAIStudioTokenCounter().acount_tokens( + **count_tokens_params_request, + ) + except litellm.APIError as error: + return provider_error_response(error, request_model, model_to_use) if result is not None: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 0e4cb856632..4c0026783f7 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1359,24 +1359,27 @@ class VertexAITokenCounter(BaseTokenCounter): from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, count_context_error_response, + gemini_count_contents, + inference_request_view, + provider_error_response, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - from litellm.llms.vertex_ai.gemini.transformation import ( - gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini - ) - count_context: Final = vertex_gemini_count_context(system=system, tools=tools) + request_view: Final = inference_request_view(messages or [], system, messages_format) + count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) resolved_contents: Final = ( contents if contents is not None - else gemini_convert_messages_with_history( - messages=messages or [] # pyright: ignore[reportArgumentType] # preserve direct caller message handling + else await gemini_count_contents( + request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request ) ) + if isinstance(resolved_contents, VertexGeminiCountContextError): + return count_context_error_response(resolved_contents, request_model, model_to_use) count_tokens_params: Final = { "model": model_to_use, @@ -1390,15 +1393,7 @@ class VertexAITokenCounter(BaseTokenCounter): **count_tokens_params_request, ) except litellm.APIError as error: - return TokenCountResponse( - total_tokens=0, - request_model=request_model, - model_used=model_to_use, - tokenizer_type="gemini", - error=True, - error_message=error.message, - status_code=error.status_code, - ) + return provider_error_response(error, request_model, model_to_use) if "totalTokens" in result: return TokenCountResponse( diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index cb0bedf0238..3f9289c8244 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -1,12 +1,20 @@ -from collections.abc import Mapping +import copy +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never from litellm.types.llms.anthropic import AnthropicMessagesTextParam -from litellm.types.utils import TokenCountResponse +from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.vertex_ai import ContentType +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse + +if TYPE_CHECKING: + from litellm.exceptions import APIError + from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult @dataclass(frozen=True, slots=True) @@ -142,3 +150,103 @@ def count_context_error_response( error_message=error.error_message, status_code=400, ) + + +def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type="gemini", + error=True, + error_message=error.message, + status_code=error.status_code, + ) + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None": + from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + chat_messages: Final = [dict(message) for message in messages] + if system is None or isinstance(system, str): + return apply_client_compaction_block_history(messages=chat_messages, system=system) + try: + blocks: Final = _SYSTEM_BLOCKS.validate_python(system) + except ValidationError: + return None + return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + + +def inference_request_view( + messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat +) -> InferenceRequestView: + """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" + from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages + + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + normalized: Final = normalize_native_messages(messages) + compacted: Final = _client_compaction(normalized, system) + if compacted is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=compacted.messages, system=compacted.system) + case _: + assert_never(messages_format) + + +def _chat_messages( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> Sequence[AllMessageValues]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + match messages_format: + case "anthropic": + return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it + model=model, + custom_llm_provider=custom_llm_provider, + ) + case "openai": + return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place + case _: + assert_never(messages_format) + + +async def gemini_count_contents( + messages: Sequence[Mapping[str, object]], + messages_format: TokenCountMessagesFormat, + model: str, + custom_llm_provider: Literal["vertex_ai", "gemini"], + litellm_params: Mapping[str, object], +) -> Sequence[ContentType] | VertexGeminiCountContextError: + from litellm.exceptions import BadRequestError + from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media + from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig + from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig() + try: + chat_messages: Final = await async_inline_remote_media( + list(_chat_messages(messages, messages_format, model, custom_llm_provider)), + should_inline=should_inline_remote_media(custom_llm_provider), + ) + return await asyncify(config.transform_messages)( + messages=chat_messages, model=model, litellm_params=dict(litellm_params) + ) + except BadRequestError as error: + return VertexGeminiCountContextError(error_message=error.message) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 57327731863..895d59965b7 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote @@ -1318,6 +1318,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool: ) +def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]: + return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + + async def async_transform_request_body( gemini_api_key: str | None, messages: list[AllMessageValues], @@ -1360,7 +1364,7 @@ async def async_transform_request_body( ) inlined_messages: Final = await async_inline_remote_media( - messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines + messages, should_inline=should_inline_remote_media(custom_llm_provider) ) if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 1a148b19f8d..7fef79845a1 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -167,6 +167,93 @@ def test_messages_and_count_tokens_send_matching_gemini_context( assert count_body == _count_tokens_request_from_generate_content(generated_body) +_LOOKUP_TOOL: Final[JsonValue] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the document."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + }, + ], + } + ], + }, + "base64-image": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "tool-turns": { + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Look up the value."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}], + }, + ], + }, + } +) + + +def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") + + +def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: + if "contents" in count_body: + return count_body["contents"] + return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES)) +def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + requests: Final = peer.drain() + generated_body: Final = _json_body( + next(request for request in requests if request.target.endswith(":generateContent")) + ) + count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens"))) + assert _counted_contents(count_body) == generated_body["contents"] + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 9a47d01ec46..e54023c6596 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -5,6 +5,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -100,6 +101,58 @@ async def test_google_ai_studio_count_tokens_rejects_image_system_without_reques assert peer.request_body is None +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.request_body is None + + +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None: + rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}' + client: Final = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request)) + ) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + + @pytest.mark.asyncio @pytest.mark.parametrize( ("system_instruction", "tools", "expected_body"), @@ -129,9 +182,7 @@ async def test_gemini_count_tokens_sends_optional_request_fields( expected_body: dict[str, JsonValue], ) -> None: peer: Final = _GeminiCountTokensPeer() - token_counter: Final = GeminiCountTokensHandler( - client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)) - ) + token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) result: Final = await token_counter.acount_tokens( contents=[{"role": "user", "parts": [{"text": "hello"}]}], @@ -177,18 +228,10 @@ class TestGeminiModelInfo: # Test edge cases where model names end with characters from "models/" # These would be incorrectly processed if using strip("models/") instead of replace("models/", "") models = [ - { - "name": "models/gemini-1.5-pro" - }, # ends with 'o' - would become "gemini-1.5-pr" with strip() - { - "name": "models/test-model" - }, # ends with 'l' - would become "gemini/test-mode" with strip() - { - "name": "models/custom-models" - }, # ends with 's' - would become "gemini/custom-model" with strip() - { - "name": "models/demo" - }, # ends with 'o' - would become "gemini/dem" with strip() + {"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip() + {"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip() + {"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip() + {"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip() ] result = gemini_model_info.process_model_name(models) @@ -239,16 +282,10 @@ class TestGoogleAIStudioTokenCounter: token_counter = GoogleAIStudioTokenCounter() # Test with gemini provider - should return True - assert ( - token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) - is True - ) + assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True # Test with other providers - should return False - assert ( - token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) - is False - ) + assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False assert token_counter.should_use_token_counting_api("anthropic") is False assert token_counter.should_use_token_counting_api("vertex_ai") is False @@ -320,9 +357,7 @@ class TestGoogleAIStudioTokenCounter: "functionResponse": { "id": "read_many_files-1757526647518-730a691aac11c", # This should be removed "name": "read_many_files", - "response": { - "output": "No files matching the criteria were found or all were skipped." - }, + "response": {"output": "No files matching the criteria were found or all were skipped."}, } } ], @@ -331,9 +366,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_with_id - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id) # Verify the 'id' field was removed function_response = cleaned_contents[1]["parts"][0]["functionResponse"] @@ -342,8 +375,7 @@ class TestGoogleAIStudioTokenCounter: assert "response" in function_response assert function_response["name"] == "read_many_files" assert ( - function_response["response"]["output"] - == "No files matching the criteria were found or all were skipped." + function_response["response"]["output"] == "No files matching the criteria were found or all were skipped." ) def test_clean_contents_for_gemini_api_preserves_other_fields(self): @@ -359,9 +391,7 @@ class TestGoogleAIStudioTokenCounter: ] # Clean the contents - cleaned_contents = token_counter._clean_contents_for_gemini_api( - contents_without_function_response - ) + cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response) # Verify the contents are unchanged assert cleaned_contents == contents_without_function_response diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 444920e1ecc..06dacedde67 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,8 +1,18 @@ -from typing import Final +import asyncio +import base64 +import time +import uuid +from collections.abc import Generator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from threading import Thread +from typing import Final, Literal +from unittest.mock import patch import pytest -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, truncate_tool_name, @@ -10,6 +20,8 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContext, VertexGeminiCountContextError, + gemini_count_contents, + inference_request_view, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -193,3 +205,302 @@ def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None: assert isinstance(context, VertexGeminiCountContextError) assert "tools" in context.error_message + + +_ANTHROPIC_CONVERSATION: Final = ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize both"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}}, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]}, +) + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + contents: Final = await gemini_count_contents( + _ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Summarize both"}, + {"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}}, + {"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}}, + ], + }, + {"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]}, + ] + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None: + contents: Final = await gemini_count_contents( + ({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}), + "openai", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + + assert contents == [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"text": "hi"}]}, + ] + + +_PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) +_DOWNLOAD_SECONDS: Final = 0.5 + + +class _SlowPngHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + time.sleep(_DOWNLOAD_SECONDS) + self.send_response(200) + self.send_header("Content-Type", "image/png") + self.send_header("Content-Length", str(len(_PNG_BYTES))) + self.end_headers() + self.wfile.write(_PNG_BYTES) + + def log_message(self, format: str, *args: object) -> None: + return None + + +@contextmanager +def _slow_png_server() -> Generator[str, None, None]: + with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server: + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + + +async def _event_loop_wake_delay(nap_seconds: float) -> float: + started: Final = time.monotonic() + await asyncio.sleep(nap_seconds) + return time.monotonic() - started + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_remote_image_downloads_without_blocking_the_event_loop( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with _slow_png_server() as base_url: + image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png" + wake_delay, contents = await asyncio.gather( + _event_loop_wake_delay(0.05), + gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it"}, + {"type": "image", "source": {"type": "url", "url": image_url}}, + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ), + ) + + assert contents == [ + { + "role": "user", + "parts": [ + {"text": "Describe it"}, + {"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}}, + ], + } + ] + assert wake_delay < _DOWNLOAD_SECONDS / 2 + + +class _MissingImageHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + return None + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +async def test_image_that_cannot_be_fetched_returns_an_error_value( + custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server: + Thread(target=server.serve_forever, daemon=True).start() + try: + contents = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png", + }, + } + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + custom_llm_provider, + {}, + ) + finally: + server.shutdown() + + assert isinstance(contents, VertexGeminiCountContextError) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + ) + + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + ) + + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = inference_request_view(messages, "Be brief.", "openai") + + assert view.messages == messages + assert view.system == "Be brief." + + +async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/png", + ) as lookup: + contents: Final = await gemini_count_contents( + ( + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}} + ], + }, + ), + "anthropic", + "gemini-2.5-flash", + "vertex_ai", + {"vertex_project": "deployment-project"}, + ) + + assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project" + assert not isinstance(contents, VertexGeminiCountContextError) + assert contents[0]["parts"][0] == { + "file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"} + } + + +@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) +def test_anthropic_replayed_web_search_result_is_counted_as_its_text( + custom_llm_provider: Literal["vertex_ai", "gemini"], +) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final = ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ) + + view: Final = inference_request_view(messages, None, "anthropic") + contents: Final = asyncio.run( + gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + ) + + assert snippet in str(contents) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 304424b2acc..8528724a24e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,5 +1,6 @@ import time from collections.abc import Iterator +from contextlib import contextmanager from typing import Final from unittest.mock import AsyncMock, patch @@ -7,6 +8,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -24,16 +26,21 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler +from litellm.types.utils import TokenCountMessagesFormat @pytest.mark.asyncio -async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> None: +@pytest.mark.parametrize("messages_format", ("openai", "anthropic")) +async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( + messages_format: TokenCountMessagesFormat, +) -> None: result: Final = await VertexAITokenCounter().count_tokens( model_to_use="gemini-4-argon", messages=[{"role": "user", "content": "Count this"}], contents=None, request_model="vertex_ai/gemini-4-argon", system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}], + messages_format=messages_format, ) assert result is not None @@ -42,6 +49,88 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> assert "text blocks" in (result.error_message or "") +class _VertexCountTokensPeer: + def __init__(self, status_code: int) -> None: + self.status_code: Final = status_code + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + if self.status_code == 200: + return httpx.Response(200, json={"totalTokens": 37}, request=request) + return httpx.Response( + self.status_code, + json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + request=request, + ) + + +@contextmanager +def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]: + peer: Final = _VertexCountTokensPeer(status_code) + with ( + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + patch.object( + vertex_count_tokens_handler.VertexAITokenCounter, + "_ensure_access_token_async", + new=AsyncMock(return_value=("fake-token", "test-project")), + ), + ): + yield peer + + +_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}} + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None: + with _vertex_count_tokens_peer(400) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "Invalid function name" in (result.error_message or "") + assert len(peer.requests) == 1 + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "user_url_validation", True) + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}], + } + ], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert "http://127.0.0.1:9/image.png" in (result.error_message or "") + assert peer.requests == () + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location From 081cfa2a6264b9a4de8b647efefa9e72a3b296f0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 19:47:12 -0700 Subject: [PATCH 03/82] fix(proxy): reset end-user spend in committed pages so the budget cascade survives a large LiteLLM_EndUserTable (#46502) * fix(proxy): reset end-user spend in committed pages so the budget cascade survives large LiteLLM_EndUserTable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): collapse _enduser_ids_by_tier return to one line for ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark the end-user page IN list as bounded for the unbounded-IN check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/reset_budget_job.py | 110 +++++-- litellm/repositories/unit_of_work.py | 17 +- .../common_utils/test_reset_budget_job.py | 276 ++++++++++++------ tests/unit/repositories/test_unit_of_work.py | 24 +- 4 files changed, 315 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 61330d12ac3..ccf2232e379 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum +from itertools import chain from types import MappingProxyType from typing import Final, Generic, Literal, Protocol, TypeVar @@ -65,6 +66,7 @@ from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, + enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, ) from litellm.repositories.verification_token_repository import ( @@ -222,8 +224,8 @@ def _budget_link_where( def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: """Customers whose cached spend a committed reset of these tiers invalidated. - Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which - post-commit would match nobody. + ``_enduser_reset_where`` adds the ``spend > 0`` filter, which post-commit + would match nobody. """ linked: Final = _budget_link_where(budget_ids) default_budget_id: Final = litellm.max_end_user_budget_id @@ -251,26 +253,44 @@ def _queue_budget_linked_resets( writes.queue_spend_zero(where=_budget_link_where(plain_ids, extra)) -def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None: - """End users reset on the budget link like every other gated table, plus a - NULL-budget_id branch: rows created implicitly persist no link and ride the - default tier (litellm.max_end_user_budget_id). +def _enduser_reset_where(cascade: "_BudgetCascade") -> Mapping[str, object]: + """End users with spend to reset: linked to an expiring tier, or carrying + no link and riding the default tier (litellm.max_end_user_budget_id).""" + return {**_enduser_invalidation_where(cascade.budget_ids), **_SPENT_ROWS_WHERE} - Matching on the link rather than enumerating user ids keeps a statement's - bind count proportional to the expiring tiers instead of the customer - population, which past ~32,700 dependents exceeds PostgreSQL's per-statement - bind ceiling and wedges the cascade permanently (#40564). + +def _enduser_tier(row: _EndUserRow) -> str | None: + return row.budget_id if row.budget_id is not None else litellm.max_end_user_budget_id + + +def _enduser_ids_by_tier(rows: Sequence[_EndUserRow]) -> Mapping[str | None, tuple[str, ...]]: + tiers: Final = tuple(dict.fromkeys(_enduser_tier(row) for row in rows)) + return MappingProxyType({tier: tuple(row.user_id for row in rows if _enduser_tier(row) == tier) for tier in tiers}) + + +def _queue_enduser_page_resets( + writes: LinkedSpendResetWrites, rows: Sequence[_EndUserRow], cascade: "_BudgetCascade" +) -> None: + """The statements of ``_queue_budget_linked_resets`` for one page of end + users, keyed by their user ids so each statement touches at most one page. + + A statement matching the budget link reads the whole end-user table (no + index on budget_id) and updates every linked row at once, which past a few + million end users outlives the prisma client's 30s request timeout (#46404). + A page of ids stays under PostgreSQL's per-statement bind ceiling (#40564) + and the spend predicates still classify each row at write time. """ - _queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE) - default_budget_id: Final = litellm.max_end_user_budget_id - if default_budget_id is None or default_budget_id not in cascade.budget_ids: - return - cap: Final = cascade.rollover_caps.get(default_budget_id) + for tier, user_ids in _enduser_ids_by_tier(rows).items(): + _queue_enduser_tier_resets(writes, user_ids, cap=cascade.rollover_caps.get(tier) if tier is not None else None) + + +def _queue_enduser_tier_resets(writes: LinkedSpendResetWrites, user_ids: Sequence[str], cap: float | None) -> None: + page: Final = {"user_id": {"in": list(user_ids)}} # bounded-ok: one page of RESET_BUDGET_JOB_BATCH_SIZE ids if cap is None: - writes.queue_spend_zero(where={"budget_id": None, **_SPENT_ROWS_WHERE}) + writes.queue_spend_zero(where={**page, **_SPENT_ROWS_WHERE}) return - writes.queue_spend_zero(where={"budget_id": None, "spend": {"gt": 0, "lte": cap}}) - writes.queue_spend_decrement(where={"budget_id": None, "spend": {"gt": cap}}, amount=cap) + writes.queue_spend_zero(where={**page, "spend": {"gt": 0, "lte": cap}}) + writes.queue_spend_decrement(where={**page, "spend": {"gt": cap}}, amount=cap) @dataclass(frozen=True, slots=True) @@ -824,6 +844,49 @@ class ResetBudgetJob: ), ) + async def _reset_enduser_spend(self, cascade: _BudgetCascade) -> int: + """Reset the end users the expiring tiers gate, one committed page at a time. + + Runs before ``_commit_budget_cascade`` so a failure here leaves the + tiers due: the next tick resumes with the pages still carrying spend. + Each page invalidates its own caches right after its commit, so no + zeroed row sits behind a cached spend that still rejects requests. + Capless like ``_invalidate_enduser_caches``, for its reason. + """ + if not cascade.budget_ids: + return 0 + where: Final = _enduser_reset_where(cascade) + walk = _EndUserWalk() # rebind-ok: keyset cursor advances one committed page at a time + while walk.cursor is not None: + walk = await self._reset_enduser_page(cascade, where=where, cursor=walk.cursor, reached=walk.invalidated) + return walk.invalidated + + async def _reset_enduser_page( + self, cascade: _BudgetCascade, where: Mapping[str, object], cursor: str, reached: int + ) -> _EndUserWalk: + rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor) + if not rows: + return _EndUserWalk(cursor=None, invalidated=reached) + await self._with_db_write_retry( + lambda: self._commit_enduser_page(rows, cascade), + reason="reset_budget_write_endusers_failure", + ) + await self._invalidate_caches( + counter_keys=tuple(_enduser_counter_key(row) for row in rows), + cache_keys=tuple(chain.from_iterable(_enduser_cache_keys(row) for row in rows)), + ) + walked: Final = reached + len(rows) + if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: + return _EndUserWalk(cursor=None, invalidated=walked) + return _EndUserWalk(cursor=rows[-1].user_id, invalidated=walked) + + async def _commit_enduser_page(self, rows: Sequence[_EndUserRow], cascade: _BudgetCascade) -> None: + async with ( + db_span("reset_budget_endusers", "LiteLLM_EndUserTable"), + enduser_spend_reset_unit_of_work(self._new_batch) as writes, + ): + _queue_enduser_page_resets(writes, rows, cascade) + async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None: """Zero the gated spend and advance ``budget_reset_at`` in one transaction. @@ -851,7 +914,6 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) - _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) @@ -878,6 +940,7 @@ class ResetBudgetJob: return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e) try: + await self._reset_enduser_spend(cascade) await self._commit_budget_cascade(cascade) except Exception as e: return _BudgetCascadeFailed(cascade=cascade, error=e) @@ -896,9 +959,10 @@ class ResetBudgetJob: """ Resets the spend a budget tier gates (end users, team members, keys, orgs, tags, model access groups) and advances the tier's - budget_reset_at, atomically. + budget_reset_at. End users commit in pages first, everything else + lands with the window advance in one transaction. - Caches are invalidated only after the transaction commits, so a failed + Caches are invalidated only after the commit they follow, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. """ await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk) @@ -929,8 +993,8 @@ class ResetBudgetJob: case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " - "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " - "next run: %s", + "group spend, plus budget_reset_at); the budgets stay due for the next run, which resumes " + "with the end-user pages not committed yet: %s", error, exc_info=error, ) diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 10f092c1cb2..95f4f5e81be 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -97,11 +97,13 @@ class SpendResetUnitOfWork: @dataclass(frozen=True, slots=True) class BudgetCascadeUnitOfWork: - """Every write a budget-tier reset performs, bound to one batch. + """Every write a budget-tier reset performs in its final batch. The dependent spend rows and the budget rows' ``budget_reset_at`` advance must land together: advancing the window without zeroing the spend it gates leaves the dependents pinned at their cap until the next window. + End users are the exception, paged and committed ahead of this batch by + ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ team_memberships: LinkedSpendResetWrites @@ -110,7 +112,6 @@ class BudgetCascadeUnitOfWork: tags: LinkedSpendResetWrites model_access_groups: LinkedSpendResetWrites projects: LinkedSpendResetWrites - endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -137,7 +138,17 @@ async def budget_cascade_unit_of_work( tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), projects=LinkedSpendResetWrites(table=batch.litellm_projecttable), - endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) await batch.commit() + + +@asynccontextmanager +async def enduser_spend_reset_unit_of_work( + new_batch: Callable[[], PrismaBatch], +) -> AsyncGenerator[LinkedSpendResetWrites, None]: + """One page of end-user spend resets, committed on its own so a statement + never grows with the end-user population.""" + batch: Final = new_batch() + yield LinkedSpendResetWrites(table=batch.litellm_endusertable) + await batch.commit() diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index c17aed49014..70f7774e691 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -435,7 +435,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): - """End-user spend is zeroed and the tier's window advances, in one batch.""" + """End-user spend is zeroed in its own committed page, then the tier's window advances.""" now = datetime.now(timezone.utc) test_budget = _budget_row(budget_id="test-budget-1", budget_duration="1d", budget_reset_at=now) @@ -446,11 +446,12 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): "spend": 20.0, "litellm_budget_table": test_budget, "user_id": "test-enduser-1", + "budget_id": "test-budget-1", }, ) mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -458,10 +459,13 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["test-enduser-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] + assert [batcher.committed for batcher in mock_prisma_client.db.batchers] == [True, True] + assert [call["table"] for call in mock_prisma_client.db.batchers[0].calls] == ["enduser"] + assert "enduser" not in {call["table"] for call in mock_prisma_client.db.batchers[1].calls} budget_writes = _batch_writes(mock_prisma_client, "budget") assert len(budget_writes) == 1 @@ -470,6 +474,112 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): assert set(budget_writes[0]["data"].keys()) == {"budget_reset_at"} +def test_enduser_reset_commits_in_pages_before_the_window_advances(reset_budget_job, mock_prisma_client): + """Regression for #46404: one statement over every end user a tier gates + reads the whole table and outlives the prisma client's 30s request timeout + once the table holds millions of rows, so the batch rolls back and + budget_reset_at never moves. The reset has to land one committed page at a + time, keyed by user id, with the window advance after the last page.""" + mock_prisma_client.data["budget"] = [_budget_row(budget_id="shared-tier", budget_duration="1d")] + population: Final = RESET_BUDGET_JOB_BATCH_SIZE * 2 + 3 + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=4.0, user_id=f"cust-{index:06d}", budget_id="shared-tier") + for index in range(population) + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + batchers: Final = mock_prisma_client.db.batchers + assert [batcher.committed for batcher in batchers] == [True, True, True, True] + pages: Final = batchers[:3] + assert [[call["table"] for call in page.calls] for page in pages] == [["enduser"]] * 3 + assert [len(page.calls[0]["where"]["user_id"]["in"]) for page in pages] == [ + RESET_BUDGET_JOB_BATCH_SIZE, + RESET_BUDGET_JOB_BATCH_SIZE, + 3, + ] + assert [page.calls[0]["where"]["spend"] for page in pages] == [{"gt": 0}] * 3 + assert [page.calls[0]["data"] for page in pages] == [{"spend": 0}] * 3 + assert pages[0].calls[0]["where"]["user_id"]["in"][0] == "cust-000000" + assert pages[2].calls[0]["where"]["user_id"]["in"] == ["cust-001000", "cust-001001", "cust-001002"] + cascade: Final = batchers[3] + assert "enduser" not in {call["table"] for call in cascade.calls} + assert [call["where"] for call in cascade.calls if call["table"] == "budget"] == [{"budget_id": "shared-tier"}] + + +def test_enduser_page_failure_leaves_the_tier_due(reset_budget_job, mock_prisma_client): + """A page that fails after earlier pages committed must not advance the + window: the committed pages are at spend 0 and drop out of the next tick's + ``spend > 0`` read, so the next tick resumes where this one stopped.""" + mock_prisma_client.data["budget"] = [_budget_row(budget_id="shared-tier", budget_duration="1d")] + endusers: Final = mock_prisma_client.db.litellm_endusertable + endusers.set_find_many_results( + [ + types.SimpleNamespace(spend=4.0, user_id=f"cust-{index:06d}", budget_id="shared-tier") + for index in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) + ] + ) + endusers.set_find_many_error(1, RuntimeError("connection reset while paging end users")) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert _batch_writes(mock_prisma_client, "budget") == [] + enduser_writes: Final = _batch_writes(mock_prisma_client, "enduser") + assert len(enduser_writes) == 1 + assert len(enduser_writes[0]["where"]["user_id"]["in"]) == RESET_BUDGET_JOB_BATCH_SIZE + + +def test_enduser_page_keeps_each_tier_s_rollover_statements_to_its_own_users( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """One page can hold end users of several expiring tiers. Each tier's + zero-then-decrement pair must name only that tier's user ids and its own + cap, or a 15.0 spend on a 20.0 tier would be carried as 5.0.""" + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [ + _budget_row(budget_id="tier-10", budget_duration="1d", max_budget=10.0), + _budget_row(budget_id="tier-20", budget_duration="1d", max_budget=20.0), + ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=15.0, user_id="cust-a", budget_id="tier-10"), + types.SimpleNamespace(spend=15.0, user_id="cust-b", budget_id="tier-20"), + types.SimpleNamespace(spend=25.0, user_id="cust-c", budget_id="tier-10"), + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert _batch_writes(mock_prisma_client, "enduser") == [ + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-a", "cust-c"]}, "spend": {"gt": 0, "lte": 10.0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-a", "cust-c"]}, "spend": {"gt": 10.0}}, + "data": {"spend": {"decrement": 10.0}}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-b"]}, "spend": {"gt": 0, "lte": 20.0}}, + "data": {"spend": 0}, + }, + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["cust-b"]}, "spend": {"gt": 20.0}}, + "data": {"spend": {"decrement": 20.0}}, + }, + ] + + def test_reset_budget_all(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime now = datetime.now(timezone.utc) @@ -529,6 +639,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): "spend": 20.0, "litellm_budget_table": test_budget, "user_id": "test-enduser-1", + "budget_id": "test-budget-1", }, ) @@ -536,7 +647,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): mock_prisma_client.data["user"] = [test_user] mock_prisma_client.data["team"] = [test_team] mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) # Run the test asyncio.run(reset_budget_job.reset_budget()) @@ -559,7 +670,7 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": ["test-budget-1"]}, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["test-enduser-1"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, } ] @@ -579,7 +690,6 @@ _LINKED_TABLE_CASES = [ ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), - ("enduser", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -620,34 +730,33 @@ def _bind_count(where: Dict[str, Any]) -> int: @pytest.mark.parametrize("population", [3, 40_000], ids=["small", "over-pg-bind-ceiling"]) def test_enduser_reset_bind_count_does_not_scale_with_population(reset_budget_job, mock_prisma_client, population): - """Regression for #40564. + """Regression for #40564 and #46404. - Enumerating every dependent user id put one bind variable per customer into - a single prepared statement. Past PostgreSQL's ceiling the statement could - not be parsed at all, so the whole atomic cascade rolled back, - budget_reset_at never advanced, and the tier stayed due on every later tick - forever. Matching on the budget link keeps the statement the same size no - matter how many customers share a tier. + Enumerating every dependent user id in one statement put one bind variable + per end user into a single prepared statement, which past PostgreSQL's + ceiling could not be parsed at all. Matching on the budget link instead made + one statement read and update the whole end-user table, which past a few + million rows outlives the prisma client's 30s request timeout. Paging by + user id keeps every statement at one page of binds however many end users + share a tier, and the statements have to cover the whole population. """ budget = _budget_row(budget_id="shared-tier", budget_duration="1d") mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - types.SimpleNamespace( - spend=1.0, - litellm_budget_table=budget, - user_id=f"cust-{index:08d}", - budget_id="shared-tier", - ) - for index in range(population) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + types.SimpleNamespace(spend=1.0, user_id=f"cust-{index:08d}", budget_id="shared-tier") + for index in range(population) + ] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) writes = _batch_writes(mock_prisma_client, "enduser") - assert [_bind_count(write["where"]) for write in writes] == [2], ( - f"the cascade must not enumerate {population} user ids: past " - f"{_POSTGRES_MAX_BIND_VARIABLES} binds PostgreSQL refuses the statement, got {writes[:1]}" - ) + assert max(_bind_count(write["where"]) for write in writes) <= RESET_BUDGET_JOB_BATCH_SIZE + 1 + assert RESET_BUDGET_JOB_BATCH_SIZE + 1 < _POSTGRES_MAX_BIND_VARIABLES + assert sorted(user_id for write in writes for user_id in write["where"]["user_id"]["in"]) == [ + f"cust-{index:08d}" for index in range(population) + ] assert _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"] is not None @@ -778,6 +887,7 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock "spend": 30.0, "litellm_budget_table": test_budget, "user_id": "enduser-explicit", + "budget_id": default_budget_id, }, ) @@ -811,40 +921,34 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock ) mock_prisma_client.data["budget"] = [test_budget] - mock_prisma_client.data["enduser"] = [enduser_with_budget] - - # Set up the DB mock for NULL-budget-id end users - mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_no_budget_row]) + mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_with_budget, enduser_no_budget_row]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Both end users are zeroed: the linked rows on the tier's budget_id, the - # implicit ones on the NULL branch that stands in for the default tier. + # Both end users are zeroed in one page: the linked row on the tier's + # budget_id and the implicit one whose NULL budget_id stands in for the + # default tier. assert _batch_writes(mock_prisma_client, "enduser") == [ { "table": "enduser", "op": "update_many", - "where": {"budget_id": {"in": [default_budget_id]}, "spend": {"gt": 0}}, - "data": {"spend": 0}, - }, - { - "table": "enduser", - "op": "update_many", - "where": {"budget_id": None, "spend": {"gt": 0}}, + "where": {"user_id": {"in": ["enduser-explicit", "enduser-implicit"]}, "spend": {"gt": 0}}, "data": {"spend": 0}, }, ] - # The post-commit invalidation walk covers both branches, so implicitly - # created customers on the default tier get their cached spend dropped too, - # and it is paged rather than reading the whole customer population. + # The paged reset and the post-commit invalidation walk both cover both + # branches, so implicitly created customers on the default tier are reset + # and get their cached spend dropped too. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls - assert len(find_many_calls) == 1 - assert find_many_calls[0]["where"]["OR"] == [ - {"budget_id": {"in": [default_budget_id]}}, - {"budget_id": None}, - ] - assert find_many_calls[0]["take"] == RESET_BUDGET_JOB_BATCH_SIZE + assert len(find_many_calls) == 2 + assert find_many_calls[0]["where"]["spend"] == {"gt": 0} + for call in find_many_calls: + assert call["where"]["OR"] == [ + {"budget_id": {"in": [default_budget_id]}}, + {"budget_id": None}, + ] + assert call["take"] == RESET_BUDGET_JOB_BATCH_SIZE litellm.max_end_user_budget_id = None @@ -879,7 +983,8 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert [call["where"] for call in find_many_calls] == [ - {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}} + {"budget_id": {"in": ["some-budget"]}, "spend": {"gt": 0}, "user_id": {"gt": ""}}, + {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}}, ] litellm.max_end_user_budget_id = None @@ -919,7 +1024,8 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert [call["where"] for call in find_many_calls] == [ - {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}} + {"budget_id": {"in": ["other-budget"]}, "spend": {"gt": 0}, "user_id": {"gt": ""}}, + {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}}, ] litellm.max_end_user_budget_id = None @@ -1671,7 +1777,8 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma on a large enough deployment rather than a slow tick. Awaiting one cache call per customer makes the last customer wait out every customer ahead of it. Both regress silently, so pin the page size, the strictly advancing cursor, - and one batched call per page. + and one batched call per page. The paged reset ahead of the commit + invalidates its pages too, so the cache calls count both walks. """ counter_cache: Final = _make_counter_invalidation_job(monkeypatch) mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] @@ -1685,7 +1792,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + all_reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + reads: Final = [read for read in all_reads if "spend" not in read["where"]] + assert len(all_reads) == 6 assert [read["take"] for read in reads] == [RESET_BUDGET_JOB_BATCH_SIZE] * 3 assert [read["where"]["user_id"]["gt"] for read in reads] == [ "", @@ -1693,8 +1802,8 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma f"cust-{RESET_BUDGET_JOB_BATCH_SIZE * 2 - 1:06d}", ] - assert counter_cache.async_delete_cache_keys.await_count == 3 - assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 3 + assert counter_cache.async_delete_cache_keys.await_count == 6 + assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6 counter_cache.async_delete_cache.assert_not_called() invalidated: Final = { @@ -1727,7 +1836,7 @@ def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_fin for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) ] ) - endusers.set_find_many_error(1, RuntimeError("connection reset while paging customers")) + endusers.set_find_many_error(3, RuntimeError("connection reset while paging customers")) logging_obj: Final = RecordingProxyLogging() job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client) @@ -1985,19 +2094,22 @@ def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monk def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mock_prisma_client, monkeypatch): - """Dependent spend and the budget_reset_at advance ride one batch.""" + """Dependent spend and the budget_reset_at advance ride one batch, after + the end-user pages have committed on their own.""" _make_counter_invalidation_job(monkeypatch) now = datetime.now(timezone.utc) budget = _budget_row(budget_id="budget-1", budget_duration="7d") mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - type("EndUser", (), {"spend": 5.0, "litellm_budget_table": budget, "user_id": "enduser-1", "budget_id": "budget-1"}) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 5.0, "user_id": "enduser-1", "budget_id": "budget-1"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - assert len(mock_prisma_client.db.batchers) == 1, "the cascade must not be split across transactions" - batcher = mock_prisma_client.db.batchers[0] + assert len(mock_prisma_client.db.batchers) == 2, "one end-user page, then the cascade" + enduser_page, batcher = mock_prisma_client.db.batchers + assert enduser_page.committed is True + assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")} assert batcher.committed is True assert {(call["table"], call["op"]) for call in batcher.calls} == { ("team_membership", "update_many"), @@ -2006,7 +2118,6 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("tag", "update_many"), ("model_access_group", "update_many"), ("project", "update_many"), - ("enduser", "update_many"), ("budget", "update_many"), } budget_write = next(call for call in batcher.calls if call["table"] == "budget") @@ -3288,13 +3399,9 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( _make_counter_invalidation_job(monkeypatch) budget = _budget_row(budget_id="budget-roll", budget_duration="1d", max_budget=10.0) mock_prisma_client.data["budget"] = [budget] - mock_prisma_client.data["enduser"] = [ - type( - "EndUser", - (), - {"spend": 15.0, "litellm_budget_table": budget, "user_id": "enduser-roll", "budget_id": "budget-roll"}, - ) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 15.0, "user_id": "enduser-roll", "budget_id": "budget-roll"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -3302,13 +3409,13 @@ def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( assert { "table": "enduser", "op": "update_many", - "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, + "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 10.0}}, "data": {"spend": {"decrement": 10.0}}, } in enduser_writes assert { "table": "enduser", "op": "update_many", - "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, + "where": {"user_id": {"in": ["enduser-roll"]}, "spend": {"gt": 0, "lte": 10.0}}, "data": {"spend": 0}, } in enduser_writes @@ -3377,13 +3484,9 @@ def test_cascade_rollover_writes_survive_sequential_execution( {"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-roll"}, ) mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership]) - mock_prisma_client.data["enduser"] = [ - type( - "EndUser", - (), - {"spend": 15.0, "litellm_budget_table": budget, "user_id": "enduser-roll", "budget_id": "budget-roll"}, - ) - ] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"spend": 15.0, "user_id": "enduser-roll", "budget_id": "budget-roll"})] + ) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -3702,8 +3805,9 @@ def _wire_cascade_reads_for_test(prisma_client, endusers=()): tables an awaitable find_many so the reads resolve instead of falling into the job's warn-and-continue path. - End users are read by the post-commit invalidation walk rather than by - ``get_data``, so callers that care about customers pass them here. + End users are read by the paged reset and the post-commit invalidation + walk rather than by ``get_data``, so callers that care about them pass + them here. """ for table in ( "litellm_teammembership", @@ -4010,6 +4114,7 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(f prisma_client.get_data.side_effect = get_data_mock prisma_client.update_data = AsyncMock() batch_calls: Final = _wire_batcher_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client, endusers=endusers) proxy_logging_obj: Final = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -4020,11 +4125,14 @@ async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(f await _await_tasks_spawned_by(job.reset_budget_for_litellm_budget_table()) - assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction" + assert prisma_client.db.batch_.call_count == 2, "one end-user page, then the cascade in one transaction" enduser_writes: Final = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == { + "user_id": {"in": ["user1", "user2", "user3", "user4", "user5", "user6"]}, + "spend": {"gt": 0}, + } assert enduser_writes[0]["data"] == {"spend": 0} budget_writes: Final = [c for c in batch_calls if c["table"] == "budget"] @@ -4098,7 +4206,7 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1 enduser_writes: Final = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}, "spend": {"gt": 0}} assert enduser_writes[0]["data"] == {"spend": 0} key_writes: Final = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"] @@ -4576,7 +4684,7 @@ async def test_service_logger_endusers_success(): enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] assert len(enduser_writes) == 1 - assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}} + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1", "user2"]}, "spend": {"gt": 0}} proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once() ( diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 9eff248b917..a08d44b9b4d 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -7,6 +7,7 @@ import pytest from litellm.repositories.unit_of_work import ( LinkedSpendResetWrites, budget_cascade_unit_of_work, + enduser_spend_reset_unit_of_work, spend_reset_unit_of_work, ) @@ -96,7 +97,6 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.tags.queue_spend_zero(where=linked) uow.model_access_groups.queue_spend_zero(where=linked) uow.projects.queue_spend_zero(where=linked) - uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -108,11 +108,31 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_tagtable.update_many", linked, {"spend": 0}), ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), ("litellm_projecttable.update_many", linked, {"spend": 0}), - ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] +async def test_enduser_page_commits_on_its_own_batch(): + """One tier can gate millions of end users, so a page of them lands in its + own transaction instead of riding the cascade's.""" + batch = FakeBatch() + + async with enduser_spend_reset_unit_of_work(lambda: batch) as writes: + writes.queue_spend_zero(where={"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}) + writes.queue_spend_decrement(where={"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, amount=5.0) + assert batch.commit_count == 0 + + assert batch.commit_count == 1 + assert batch.calls == [ + ( + "litellm_endusertable.update_many", + {"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}}, + {"spend": 0}, + ), + ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, {"spend": {"decrement": 5.0}}), + ] + + async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): """A tier deleted between the read and the commit must not abort the batch: ``update`` raises P2025 on a missing row and takes every other write in the From 34d724c609b22ec4136a1eaea764e1b2d31534a2 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 20:10:17 -0700 Subject: [PATCH 04/82] test(e2e): perf datastore limits count the Postgres statements and Redis commands a request waits on, background work reported on its own (#46798) * test(e2e): perf datastore limits count what a request waits on, background Postgres and Redis work reported on its own postgres_calls_per_request and redis_calls_per_request were whole-server totals (pg_stat_database commits, Redis INFO commands), so spend flushes and config reloads moved them and a new per-request lookup could hide in the noise. They now count the statements and commands a request waits on before it is forwarded (auth and routing reads, the rate-limiter script), from pg_stat_statements read before and after the run and a Redis MONITOR connection held during it. Spend writes, daily aggregates, config reloads and cache invalidations are postgres_background_calls_per_request and redis_background_calls_per_request. The run output lists the top statements and commands per bucket, and the counters are read right after the run so the test's own spend sample is not counted * test(e2e): perf datastore reads only the runtime user's pg_stat_statements rows, other users' rows have a NULL queryid * test(e2e): perf Redis classifier strips the litellm.caching: namespace the deployed gateway puts on every key * test(e2e): small_bodies_5rps limits from ten split-metric runs (postgres awaited 1.8, redis awaited 18, postgres background 12.5), spend buffers and cron locks classified as background * test(e2e): small_bodies_5rps asserts only the awaited Postgres and Redis counts, background lines stay printed * test(e2e): perf Redis classifier puts the cron lock script's GET and DEL in background, not awaited --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/AGENTS.md | 2 +- tests/e2e/perf/cases/small_bodies_5rps.py | 10 +- tests/e2e/perf/setup/datastore.py | 233 ++++++++++++++++-- tests/e2e/perf/setup/observe.py | 70 +++++- tests/e2e/perf/setup/performance_test_case.py | 2 + tests/e2e/perf/test_perf_cases.py | 13 +- 6 files changed, 295 insertions(+), 35 deletions(-) diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index e7a26c12db8..d267059cdc7 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres transactions and Redis commands each request cost (`postgres_calls_per_request`, `redis_calls_per_request`, from `pg_stat_database` and Redis `INFO` read by `datastore.py` before and after the run). Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index 52fe41623e3..00a9fdefdb5 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -64,11 +64,11 @@ CASE: Final = PerformanceTestCase( ), limits=( Limit("unexpected_status_count", max=3), - Limit("gateway_added_seconds", percentile=50, max=0.38), - Limit("time_to_first_token_seconds", percentile=99, max=3.0), + Limit("gateway_added_seconds", percentile=50, max=0.29), + Limit("time_to_first_token_seconds", percentile=99, max=3.7), Limit("memory_growth_megabytes", max=100), - Limit("postgres_calls_per_request", max=13), - Limit("redis_calls_per_request", max=35), - Limit("cpu_seconds_per_request", max=0.19), + Limit("postgres_calls_per_request", max=1.8), + Limit("redis_calls_per_request", max=18), + Limit("cpu_seconds_per_request", max=0.29), ), ) diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py index e123d129cc5..b4ad10b5a5a 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -1,26 +1,90 @@ -"""Postgres and Redis work counters, read before and after the traffic phase so a case can put a ceiling -on how much database and cache work the gateway does per request. Both come from the servers' own -statistics (`pg_stat_database`, Redis `INFO`), so nothing here depends on how the gateway talks to them. -`DATABASE_URL` and `REDIS_HOST` are the runner's own env, as in the other suites; without them the -counters are None and any Limit on them reports as missing.""" +"""Postgres and Redis work the gateway did during the traffic phase, split into what a request waits on and +what runs after the response. Both come from the servers' own accounting, so nothing here depends on how the +gateway talks to them: Postgres statement counts from `pg_stat_statements` (read before and after the run) and +Redis commands from a `MONITOR` connection held open for the run. `DATABASE_URL` and `REDIS_HOST` are the +runner's own env, as in the other suites; without them the counters are None and any Limit on them reports as +missing. `E2E_DATABASE_ADMIN_URL`, when set, is used once to create the `pg_stat_statements` extension. Only the +runtime user's own statements are read: other users' rows come back with a NULL queryid and hidden text.""" from __future__ import annotations import os +import re +import threading +from collections import Counter +from collections.abc import Generator +from contextlib import contextmanager from dataclasses import dataclass -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import TypeAdapter -_DB_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" -_DB_ROW: Final = TypeAdapter(tuple[int]) +if TYPE_CHECKING: + import redis + +Bucket = Literal["awaited", "background", "plumbing", "unclassified"] + +_COMMITS_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" +_STATEMENTS_QUERY: Final = ( + "SELECT queryid, calls, query FROM pg_stat_statements " + "WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND userid = (SELECT usesysid FROM pg_user WHERE usename = current_user) AND queryid IS NOT NULL" +) +_CREATE_EXTENSION: Final = "CREATE EXTENSION IF NOT EXISTS pg_stat_statements" +_COMMITS_ROW: Final = TypeAdapter(tuple[int]) +_STATEMENT_ROWS: Final = TypeAdapter(list[tuple[int, int, str]]) _INFO: Final = TypeAdapter(dict[str, object]) _INT: Final = TypeAdapter(int) +_MONITOR_LINE: Final = TypeAdapter(dict[str, object]) + +_TABLE: Final = re.compile(r'"(LiteLLM_\w+)"') +_AWAITED_TABLES: Final = frozenset( + { + "LiteLLM_VerificationToken", + "LiteLLM_UserTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembership", + "LiteLLM_EndUserTable", + "LiteLLM_BudgetTable", + "LiteLLM_ObjectPermissionTable", + } +) +_PLUMBING_VERBS: Final = frozenset({"BEGIN", "COMMIT", "ROLLBACK", "DEALLOCATE", "SET", "SHOW"}) +_REDIS_READS: Final = frozenset({"GET", "MGET", "HGET", "HGETALL", "EXISTS", "TTL", "PTTL"}) +_REDIS_NAMESPACE: Final = "litellm.caching:" +_REDIS_WRITES: Final = frozenset( + {"SET", "SETEX", "DEL", "EXPIRE", "INCRBYFLOAT", "INCR", "RPUSH", "LPUSH", "LPOP", "RPOP", "LTRIM", "HSET", "SADD"} +) + + +@dataclass(frozen=True, slots=True) +class StatementCalls: + queryid: int + query: str + calls: int + + @property + def bucket(self) -> Bucket: + """A SELECT on the key, user, team, org, end-user or budget tables is what a request waits on at auth + and routing. Transaction control and session statements are plumbing. Everything else (spend and + audit writes, daily aggregates, config and model reloads) runs after the response or on a timer.""" + verb: Final = self.query.lstrip("(").split(maxsplit=1)[0].upper() if self.query.strip() else "" + tables: Final = _TABLE.findall(self.query) + if verb in _PLUMBING_VERBS or not tables: + return "plumbing" + return "awaited" if verb == "SELECT" and tables[0] in _AWAITED_TABLES else "background" + + @property + def summary(self) -> str: + return re.sub(r"\s+", " ", self.query)[:110] @dataclass(frozen=True, slots=True) class DbCounters: commits: int + statements: tuple[StatementCalls, ...] | None @dataclass(frozen=True, slots=True) @@ -29,36 +93,169 @@ class RedisCounters: connected_clients: int +@dataclass(frozen=True, slots=True) +class RedisCommand: + command: str + key: str + inside_script: bool + count: int + + @property + def bucket(self) -> Bucket: + """Budget reads, cooldown reads and the rate-limiter script (with every command it issues) happen + before the request is forwarded. Spend increments, token counters, config rewrites, cache + invalidations (PUBLISH, DEL), the spend buffers the pods flush and any other write happen after the + response or on a timer. Reads on keys these rules do not know stay unclassified and count in neither.""" + key: Final = self.key.removeprefix(_REDIS_NAMESPACE) + if key.startswith(("litellm_config:", "daily_", "cronjob_lock:", "litellm_spend", "litellm_daily")): + return "background" + if self.inside_script or self.command in {"EVALSHA", "EVAL"}: + return "awaited" + if self.command in {"PUBLISH", "DEL", "SCRIPT", "INFO", "MONITOR"}: + return "background" + if key.startswith("spend:"): + return "awaited" if self.command in _REDIS_READS else "background" + if key.startswith("deployment:") or key.endswith(("}:requests", "}:window")): + return "awaited" + if key.endswith("}:tokens") or self.command in _REDIS_WRITES: + return "background" + return "unclassified" + + @property + def summary(self) -> str: + key: Final = re.sub(r"[0-9a-f]{16,}|sk-[A-Za-z0-9_-]+|[0-9a-f-]{36}", "*", self.key)[:70] + return f"{self.command} {key}{' (in script)' if self.inside_script else ''}".rstrip() + + +def statement_deltas(before: DbCounters | None, after: DbCounters | None) -> tuple[StatementCalls, ...] | None: + """Calls per statement during the run: `after` minus `before` by queryid, new statements whole.""" + if before is None or after is None or before.statements is None or after.statements is None: + return None + earlier: Final = {statement.queryid: statement.calls for statement in before.statements} + return tuple( + StatementCalls( + queryid=statement.queryid, query=statement.query, calls=statement.calls - earlier.get(statement.queryid, 0) + ) + for statement in after.statements + if statement.calls > earlier.get(statement.queryid, 0) + ) + + +def calls_in(bucket: Bucket, items: tuple[StatementCalls, ...] | tuple[RedisCommand, ...] | None) -> int | None: + if items is None: + return None + return sum(_count(item) for item in items if item.bucket == bucket) + + +def _count(item: StatementCalls | RedisCommand) -> int: + return item.calls if isinstance(item, StatementCalls) else item.count + + def read_db_counters() -> DbCounters | None: url: Final = os.environ.get("DATABASE_URL") if url is None: return None import psycopg + _ensure_statements_extension(url) with psycopg.connect(url) as conn: - row: Final = conn.execute(_DB_QUERY).fetchone() - (commits,) = _DB_ROW.validate_python(row) - return DbCounters(commits=commits) + (commits,) = _COMMITS_ROW.validate_python(conn.execute(_COMMITS_QUERY).fetchone()) + try: + rows = _STATEMENT_ROWS.validate_python(conn.execute(_STATEMENTS_QUERY).fetchall()) + except psycopg.errors.UndefinedTable: + return DbCounters(commits=commits, statements=None) + return DbCounters( + commits=commits, + statements=tuple(StatementCalls(queryid=queryid, query=query, calls=calls) for queryid, calls, query in rows), + ) -def read_redis_counters() -> RedisCounters | None: - """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" - host: Final = os.environ.get("REDIS_HOST") - if host is None: - return None +def _ensure_statements_extension(url: str) -> None: + """Creating the extension needs a superuser, so prefer the admin URL; the runtime user's attempt is + allowed to fail (a database without `shared_preload_libraries = pg_stat_statements` cannot have it).""" + import psycopg + + try: + with psycopg.connect(os.environ.get("E2E_DATABASE_ADMIN_URL", url), autocommit=True) as conn: + _ = conn.execute(_CREATE_EXTENSION) + except psycopg.Error: + return + + +def _redis_client() -> redis.Redis[str]: import redis - client: Final = redis.Redis( - host=host, + return redis.Redis( + host=os.environ["REDIS_HOST"], port=int(os.environ.get("REDIS_PORT", "6379")), ssl=os.environ.get("REDIS_SSL", "false").lower() == "true", password=os.environ.get("REDIS_PASSWORD") or None, decode_responses=True, socket_connect_timeout=2, ) + + +def read_redis_counters() -> RedisCounters | None: + """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" + if os.environ.get("REDIS_HOST") is None: + return None + client: Final = _redis_client() stats: Final = _INFO.validate_python(client.info("stats")) clients: Final = _INFO.validate_python(client.info("clients")) return RedisCounters( commands=_INT.validate_python(stats["total_commands_processed"]), connected_clients=_INT.validate_python(clients["connected_clients"]), ) + + +class RedisMonitor: + """Every command Redis executed while the monitor was open, counted by command, key and whether a Lua + script issued it. `MONITOR` echoes each command to this connection, so it costs Redis one write per + command; fine at the rates these cases run.""" + + def __init__(self) -> None: + self._counts: Counter[tuple[str, str, bool]] = Counter() + self._thread: threading.Thread | None = None + self._pool: redis.ConnectionPool | None = None + self.commands: tuple[RedisCommand, ...] | None = None + + def _run(self) -> None: + import redis + + client: Final = _redis_client() + self._pool = client.connection_pool # pyright: ignore[reportAny] # redis.Redis leaves connection_pool untyped + with client.monitor() as monitor: + try: + while True: + line = _MONITOR_LINE.validate_python(monitor.next_command()) + words = str(line["command"]).split(" ", 2) + key = words[1].strip('"') if len(words) > 1 else "" + self._counts[(words[0].upper(), key, line["client_type"] == "lua")] += 1 + except (redis.ConnectionError, OSError, ValueError): + return + + def start(self) -> None: + self._thread = threading.Thread(target=self._run, name="redis-monitor", daemon=True) + self._thread.start() + + def stop(self) -> tuple[RedisCommand, ...]: + if self._pool is not None: + self._pool.disconnect() + if self._thread is not None: + self._thread.join(timeout=5) + return tuple( + RedisCommand(command=command, key=key, inside_script=inside_script, count=count) + for (command, key, inside_script), count in self._counts.items() + ) + + +@contextmanager +def redis_monitor() -> Generator[RedisMonitor, None, None]: + """Counts Redis commands for the duration of the block; `.commands` is set once the block ends.""" + monitor: Final = RedisMonitor() + if os.environ.get("REDIS_HOST") is not None: + monitor.start() + try: + yield monitor + finally: + monitor.commands = monitor.stop() if os.environ.get("REDIS_HOST") is not None else None diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index a322e5a7443..98391a02418 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -15,7 +15,14 @@ from e2e_http import NoBody, get_external, post_external, unwrap from memory_readings import RssCapture, RssReading, WorkerKey from pydantic import BaseModel -from setup.datastore import DbCounters, RedisCounters +from setup.datastore import ( + DbCounters, + RedisCommand, + RedisCounters, + StatementCalls, + calls_in, + statement_deltas, +) from setup.mocks.reply_common import Provider from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase from setup.wire import WireRecord @@ -106,6 +113,7 @@ class Observation: db_after: DbCounters | None = None redis_before: RedisCounters | None = None redis_after: RedisCounters | None = None + redis_commands: tuple[RedisCommand, ...] | None = None gateway_cpu_seconds: float | None = None def per_request(self, before: int | None, after: int | None) -> float | None: @@ -113,6 +121,13 @@ class Observation: return None return (after - before) / len(self.records) + def count_per_request(self, count: int | None) -> float | None: + return None if count is None or not self.records else count / len(self.records) + + @property + def statements(self) -> tuple[StatementCalls, ...] | None: + return statement_deltas(self.db_before, self.db_after) + @property def retried_markers(self) -> frozenset[str]: counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None) @@ -191,11 +206,13 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio if metric == "unexpected_status_count": return float(observation.unexpected_status_count()) if metric == "postgres_calls_per_request": - db_before, db_after = observation.db_before, observation.db_after - return observation.per_request(db_before and db_before.commits, db_after and db_after.commits) + return observation.count_per_request(calls_in("awaited", observation.statements)) + if metric == "postgres_background_calls_per_request": + return observation.count_per_request(calls_in("background", observation.statements)) if metric == "redis_calls_per_request": - redis_before, redis_after = observation.redis_before, observation.redis_after - return observation.per_request(redis_before and redis_before.commands, redis_after and redis_after.commands) + return observation.count_per_request(calls_in("awaited", observation.redis_commands)) + if metric == "redis_background_calls_per_request": + return observation.count_per_request(calls_in("background", observation.redis_commands)) if metric == "cpu_seconds_per_request": cpu: Final = observation.gateway_cpu_seconds return None if cpu is None or not observation.records else cpu / len(observation.records) @@ -234,7 +251,9 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( ("time_to_first_token_seconds", 99), ("memory_growth_megabytes", None), ("postgres_calls_per_request", None), + ("postgres_background_calls_per_request", None), ("redis_calls_per_request", None), + ("redis_background_calls_per_request", None), ("cpu_seconds_per_request", None), ) @@ -298,6 +317,7 @@ def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) flag = " OVER" if row.over else " MISSING" if row.missing else "" lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}") lines.extend(_diagnostics(observation)) + lines.extend(_datastore_lines(observation)) for record in failures: lines.append( f" wrong: seq {record.seq} {record.caller} {record.endpoint}{' stream' if record.stream else ''} status {record.status}: {record.error}" @@ -316,9 +336,49 @@ def _diagnostics(observation: Observation) -> list[str]: f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", f" retried share {len(observation.retried_markers) / count:.4f}", f" spend logged share {'n/a' if share is None else f'{share:.3f}'}", + f" postgres transactions per request {_fmt(observation.per_request(_commits(observation.db_before), _commits(observation.db_after)))}", + f" redis commands per request {_fmt(observation.per_request(_commands(observation.redis_before), _commands(observation.redis_after)))}", f" redis connected clients {_clients(observation.redis_before)} -> {_clients(observation.redis_after)}", ] +def _datastore_lines(observation: Observation) -> list[str]: + """What the per-request Postgres and Redis counts are made of, so a red run names the statement or + command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit.""" + count: Final = max(1, len(observation.records)) + lines: list[str] = [] + for title, items in ( + ("postgres statements", observation.statements), + ("redis commands", observation.redis_commands), + ): + if items is None: + lines.append(f" {title}: n/a") + continue + for bucket in ("awaited", "background", "plumbing", "unclassified"): + in_bucket = sorted((item for item in items if item.bucket == bucket), key=_item_count, reverse=True) + total = sum(_item_count(item) for item in in_bucket) + if not in_bucket: + continue + lines.append(f" {title} {bucket}: {total / count:.3f} per request, {len(in_bucket)} distinct") + lines.extend(f" {_item_count(item) / count:8.3f} {item.summary}" for item in in_bucket[:8]) + return lines + + +def _item_count(item: StatementCalls | RedisCommand) -> int: + return item.calls if isinstance(item, StatementCalls) else item.count + + +def _commits(counters: DbCounters | None) -> int | None: + return None if counters is None else counters.commits + + +def _commands(counters: RedisCounters | None) -> int | None: + return None if counters is None else counters.commands + + +def _fmt(value: float | None) -> str: + return "n/a" if value is None else f"{value:.3f}" + + def _clients(counters: RedisCounters | None) -> str: return "n/a" if counters is None else str(counters.connected_clients) diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index d90b262329c..76e6efa3203 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -22,8 +22,10 @@ Metric = Literal[ "time_to_first_token_seconds", "memory_growth_megabytes", "postgres_calls_per_request", + "postgres_background_calls_per_request", "cpu_seconds_per_request", "redis_calls_per_request", + "redis_background_calls_per_request", ] diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index d077cb9f8c7..d89a8634ab3 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -9,11 +9,13 @@ import time from typing import Final import pytest +from cases import CASES from idp import Keycloak from lifecycle import ResourceManager from memory_readings import read_rss_everywhere from models import SpendLogsParams -from cases import CASES +from proxy_client import ProxyClient +from setup.datastore import read_db_counters, read_redis_counters, redis_monitor from setup.observe import ( Observation, evaluate, @@ -26,11 +28,9 @@ from setup.observe import ( ) from setup.performance_test_case import PerformanceTestCase from setup.plan import build_plan -from setup.datastore import read_db_counters, read_redis_counters from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url from setup.run import run_plan -from proxy_client import ProxyClient pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] _SPEND_SAMPLE: Final = 40 @@ -68,9 +68,11 @@ def test_perf_case( db_before: Final = read_db_counters() redis_before: Final = read_redis_counters() run_started: Final = time.monotonic() - with gateway_cpu_sampler() as cpu: + with gateway_cpu_sampler() as cpu, redis_monitor() as monitor: outcome: Final = run_plan(case, plan, setup) run_finished: Final = time.monotonic() + db_after: Final = read_db_counters() + redis_after: Final = read_redis_counters() rss_after: Final = read_rss_everywhere(proxy) rss_read: Final = time.monotonic() assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}" @@ -81,8 +83,6 @@ def test_perf_case( proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE] ) spend_read: Final = time.monotonic() - db_after: Final = read_db_counters() - redis_after: Final = read_redis_counters() mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url)) guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else () mocks_read: Final = time.monotonic() @@ -98,6 +98,7 @@ def test_perf_case( db_after=db_after, redis_before=redis_before, redis_after=redis_after, + redis_commands=monitor.commands, gateway_cpu_seconds=cpu.seconds, ) rows, exceeded_limits = evaluate(case, observation) From 91ae7d1b7da269e05a8ad327139b267ec95993a5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:18:17 -0700 Subject: [PATCH 05/82] fix(cost): bill Responses API image generation and usage reported web search as tool usage cost (#41877) * fix(cost): bill Responses API image_generation_call tool usage Price completed image_generation_call output items from tool_usage.image_gen token counts when present, otherwise per image from the quality and size on the item, with the image model taken from tools[].model (default gpt-image-1). Add the result into tool_usage_cost alongside web search and file search, and read tool_usage from the top level or usage.tool_usage. Squashed and rebased onto main * style(cost): annotate read-only views and mark casts for the type discipline gate * fix(cost): price Responses image tool from tools[].model only, no default model or provider * refactor(cost): drop redundant status check, unread usage fields and unused output_type literal * test(cost): integration case for Responses image_generation tool token billing --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_cost_calc/tool_call_cost_tracking.py | 162 ++++++- litellm/types/llms/openai.py | 16 + .../cost_calculation/cases/openai.py | 82 ++++ .../test_tool_call_cost_tracking.py | 416 ++++++++++++++++++ 4 files changed, 654 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 6c7760d7da8..ce5dceb2e5c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -2,12 +2,13 @@ Helper utilities for tracking the cost of built-in tools. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final, Literal, cast from pydantic import ValidationError import litellm +from litellm._logging import verbose_logger from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS from litellm.litellm_core_utils.llm_cost_calc.utils import ( get_web_search_requests_from_usage, @@ -30,20 +31,36 @@ from litellm.types.utils import ( ) +def _output_item_field(output_item: object, field: str) -> object: + if isinstance(output_item, dict): + fields: Final[Mapping[str, object]] = cast(Mapping[str, object], output_item) # cast-ok: narrowed by isinstance + return fields.get(field) + return getattr(output_item, field, None) + + def _output_item_type(output_item: object) -> str | None: - item_type: Final = output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None) + item_type: Final = _output_item_field(output_item, "type") return item_type if isinstance(item_type, str) else None -def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: - tool_usage: Final = getattr(response_object, "tool_usage", None) - if tool_usage is None: +def _responses_tool_usage(response_object: ResponsesAPIResponse) -> ResponsesToolUsage | None: + top_level: Final = getattr(response_object, "tool_usage", None) + raw: Final = ( + top_level if top_level is not None else getattr(getattr(response_object, "usage", None), "tool_usage", None) + ) + if raw is None: return None try: - web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search + return ResponsesToolUsage.model_validate(raw) except ValidationError: return None - return None if web_search is None else web_search.num_requests + + +def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: + tool_usage: Final = _responses_tool_usage(response_object) + if tool_usage is None: + return None + return None if tool_usage.web_search is None else tool_usage.web_search.num_requests def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: @@ -87,31 +104,48 @@ class StandardBuiltInToolCostTracking: usage=usage, ) + image_generation_cost: Final = StandardBuiltInToolCostTracking._handle_image_generation_cost( + response_object=response_object, + custom_llm_provider=custom_llm_provider, + ) + # Handle web search if StandardBuiltInToolCostTracking.response_object_includes_web_search_call( response_object=response_object, usage=usage ): - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_web_search_cost( - model=model, - custom_llm_provider=custom_llm_provider, - usage=usage, - standard_built_in_tools_params=standard_built_in_tools_params, - response_object=response_object, + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_web_search_cost( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + standard_built_in_tools_params=standard_built_in_tools_params, + response_object=response_object, + ) ) # Handle file search if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response_object): - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_file_search_cost( + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_file_search_cost( + model=model, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=standard_built_in_tools_params, + ) + ) + + # Handle Azure assistant features + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_azure_assistant_costs( model=model, custom_llm_provider=custom_llm_provider, standard_built_in_tools_params=standard_built_in_tools_params, ) - - # Handle Azure assistant features - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_azure_assistant_costs( - model=model, - custom_llm_provider=custom_llm_provider, - standard_built_in_tools_params=standard_built_in_tools_params, ) @staticmethod @@ -210,6 +244,87 @@ class StandardBuiltInToolCostTracking: ) return max(count, 1) + @staticmethod + def _image_generation_tool_model(response_object: ResponsesAPIResponse) -> str | None: + tools: Final[Sequence[object]] = cast(Sequence[object], response_object.tools or ()) # cast-ok: dict tool union + for tool in tools: + if _output_item_field(tool, "type") != "image_generation": + continue + if isinstance(model := _output_item_field(tool, "model"), str) and model: + return model + return None + + @staticmethod + def _image_generation_call_cost(output_item: object, tool_model: str, custom_llm_provider: str | None) -> float: + from litellm.cost_calculator import ( + default_image_cost_calculator, # pyright: ignore[reportUnknownVariableType] # optional_params param is untyped + ) + + quality: Final = _output_item_field(output_item, "quality") + size: Final = _output_item_field(output_item, "size") + try: + return default_image_cost_calculator( + model=tool_model, + custom_llm_provider=custom_llm_provider, + quality=quality if isinstance(quality, str) and quality != "auto" else None, + n=1, + size=size if isinstance(size, str) and size != "auto" else None, + ) + except Exception as e: # noqa: BLE001 # pricing helpers raise bare Exception for unmapped models; bill 0.0 + verbose_logger.debug("Could not price Responses API image_generation_call item: %s", e) + return 0.0 + + @staticmethod + def _image_generation_token_cost( + response_object: ResponsesAPIResponse, tool_model: str, custom_llm_provider: str | None + ) -> float | None: + tool_usage: Final = _responses_tool_usage(response_object) + if tool_usage is None or tool_usage.image_gen is None or tool_usage.image_gen.total_tokens <= 0: + return None + try: + model_info: Final = litellm.get_model_info(model=tool_model, custom_llm_provider=custom_llm_provider) + except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; fall back + verbose_logger.debug("Could not resolve pricing for image tool model %s: %s", tool_model, e) + return None + image_gen: Final = tool_usage.image_gen + input_details: Final = image_gen.input_tokens_details + output_details: Final = image_gen.output_tokens_details + token_cost: Final = ( + (input_details.text_tokens if input_details else 0) * (model_info.get("input_cost_per_token") or 0) + + (input_details.image_tokens if input_details else 0) * (model_info.get("input_cost_per_image_token") or 0) + + (output_details.image_tokens if output_details else 0) + * (model_info.get("output_cost_per_image_token") or 0) + + (output_details.text_tokens if output_details else 0) * (model_info.get("output_cost_per_token") or 0) + ) + return token_cost if token_cost > 0 else None + + @staticmethod + def _handle_image_generation_cost(response_object: object, custom_llm_provider: str | None) -> float: + if not isinstance(response_object, ResponsesAPIResponse): + return 0.0 + output: Final[Sequence[object]] = cast(Sequence[object], response_object.output) # cast-ok: isinstance-narrowed + completed_items: Final = tuple( + output_item + for output_item in output + if _output_item_type(output_item) == "image_generation_call" + and _output_item_field(output_item, "status") == "completed" + ) + if not completed_items: + return 0.0 + tool_model: Final = StandardBuiltInToolCostTracking._image_generation_tool_model(response_object) + if tool_model is None: + verbose_logger.debug("Responses API image_generation_call has no tools[].model to price, billing 0.0") + return 0.0 + token_cost: Final = StandardBuiltInToolCostTracking._image_generation_token_cost( + response_object, tool_model, custom_llm_provider + ) + if token_cost is not None: + return token_cost + return sum( + StandardBuiltInToolCostTracking._image_generation_call_cost(output_item, tool_model, custom_llm_provider) + for output_item in completed_items + ) + @staticmethod def _handle_file_search_cost( model: str, @@ -445,9 +560,12 @@ class StandardBuiltInToolCostTracking: return False elif isinstance(response_object, ResponsesAPIResponse): # response api explicitly includes web_search_call in the output - return StandardBuiltInToolCostTracking.response_includes_output_type( + if StandardBuiltInToolCostTracking.response_includes_output_type( response_object=response_object, output_type="web_search_call" - ) + ): + return True + reported: Final = _reported_web_search_requests(response_object) + return reported is not None and reported > 0 elif usage is not None: if get_web_search_requests_from_usage(usage) is not None or ( hasattr(usage, "prompt_tokens_details") diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8731bc2905a..8cc9e912e80 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1421,10 +1421,26 @@ class WebSearchToolUsage(LiteLLMBaseModel): num_requests: NonNegativeInt +class ImageGenTokenDetails(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + image_tokens: NonNegativeInt = 0 + text_tokens: NonNegativeInt = 0 + + +class ImageGenToolUsage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + total_tokens: NonNegativeInt = 0 + input_tokens_details: ImageGenTokenDetails | None = None + output_tokens_details: ImageGenTokenDetails | None = None + + class ResponsesToolUsage(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) web_search: WebSearchToolUsage | None = None + image_gen: ImageGenToolUsage | None = None ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 9a3885af259..cfb7fda4d71 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -11192,6 +11192,87 @@ GPT_5_6_SLASH_ALIAS_STREAM: Final = GPT_5_6_STREAM.model_copy( update={"name": "gpt-5.6-slash_alias_stream", "alias": "openai/custom/gpt-5.6-$MARKER"} ) +GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS: Final = CostTrackingTestCase( + name="gpt-5.6-responses_image_generation_tool_tokens", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "draw a duck", + "tools": [{"type": "image_generation", "model": "gpt-image-next", "quality": "low", "size": "1024x1024"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "tools": [ + { + "type": "image_generation", + "background": "auto", + "model": "gpt-image-next", + "moderation": "auto", + "n": 1, + "output_compression": 100, + "output_format": "png", + "quality": "low", + "size": "1024x1024", + } + ], + "output": [ + { + "type": "image_generation_call", + "id": "ig_$REQUEST_ID", + "status": "completed", + "action": "generate", + "background": "opaque", + "output_format": "png", + "quality": "low", + "revised_prompt": "a duck", + "size": "1024x1024", + "result": "AAAA", + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + }, + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "tool_usage": { + "image_gen": { + "input_tokens": 32, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 32}, + "output_tokens": 196, + "output_tokens_details": {"image_tokens": 196, "text_tokens": 0}, + "total_tokens": 228, + }, + "web_search": {"num_requests": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.01004232, + input_cost=0.00322, + output_cost=0.005768, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.00105432, + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, @@ -11336,6 +11417,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_MIXED_ACTIONS, GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_OPEN_PAGE_ONLY, GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_MIXED_ACTIONS_STREAM, + GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS, GPT_5_3_CODEX_RESPONSES_FILE_SEARCH, GPT_5_6_RESPONSES_SERVICE_TIER_FLEX, GPT_5_6_RESPONSES_SERVICE_TIER_PRIORITY, diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 320a4eff371..c8dbbd00d66 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -593,6 +593,422 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +def _openai_responses_response(model, output, usage=None, tools=None): + return ResponsesAPIResponse.model_validate( + { + "id": "resp_1", + "created_at": 1754900000, + "model": model, + "object": "response", + "status": "completed", + "output": output, + "usage": usage or {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + **({"tools": tools} if tools is not None else {}), + } + ) + + +_ASSISTANT_MESSAGE_OUTPUT_ITEM = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], +} + +_GPT_IMAGE_1_HIGH_1024_COST_KEY = "high/1024-x-1024/gpt-image-1" +_GPT_IMAGE_1_TOOLS = [{"type": "image_generation", "model": "gpt-image-1"}] + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", None]) +def test_responses_image_generation_call_billed_as_tool_usage_cost(local_model_cost_map, custom_llm_provider): + """A completed image_generation_call bills at the tools[].model rate for its quality/size, provider given or inferred.""" + expected_image_cost = litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=None, + ) + + assert cost > 0 + assert cost == pytest.approx(expected_image_cost) + + +def test_responses_image_generation_call_without_tool_model_bills_zero(local_model_cost_map): + """With no tools[].model to price against, the image tool part bills $0 instead of guessing a model or raising.""" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + ], + usage=_image_gen_token_usage(10, 0, 50, 0), + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +def test_responses_web_search_and_image_generation_costs_are_additive(local_model_cost_map): + """A response billed for web search must still also bill its image_generation_call items.""" + model = "gpt-4o-search-preview" + image_cost = litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + image_item = { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + web_search_item = {"type": "web_search_call", "id": "ws_1", "status": "completed"} + + combined = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [web_search_item, image_item], tools=_GPT_IMAGE_1_TOOLS), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + web_search_only = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [web_search_item]), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + image_only = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [image_item], tools=_GPT_IMAGE_1_TOOLS), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert image_only == pytest.approx(image_cost) + assert web_search_only > 0 + assert combined == pytest.approx(web_search_only + image_only) + + +def test_responses_incomplete_image_generation_call_not_billed(local_model_cost_map): + """A failed image_generation_call produced no billable image, so it must cost $0.""" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "failed", + "quality": "high", + "size": "1024x1024", + "result": None, + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +def test_completion_cost_includes_responses_image_generation_tool_cost(local_model_cost_map): + """The image tool fee must flow through completion_cost on top of the token-only baseline.""" + image_item = { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + response_with_image = _openai_responses_response( + "gpt-5", [image_item, dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], tools=_GPT_IMAGE_1_TOOLS + ) + response_without_image = _openai_responses_response("gpt-5", [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)]) + + cost_with_image = litellm.completion_cost( + completion_response=response_with_image, + model="gpt-5", + custom_llm_provider="openai", + call_type="aresponses", + ) + cost_without_image = litellm.completion_cost( + completion_response=response_without_image, + model="gpt-5", + custom_llm_provider="openai", + call_type="aresponses", + ) + + assert cost_with_image > cost_without_image + assert cost_with_image - cost_without_image == pytest.approx( + litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + ) + + +def test_responses_usage_tool_usage_web_search_billed_without_output_item(local_model_cost_map): + """usage.tool_usage.web_search.num_requests bills web search even when no web_search_call item is present.""" + model = "gpt-5.4-mini" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] + + for num_requests in (1, 2): + response = _openai_responses_response( + model, + [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"web_search": {"num_requests": num_requests}}, + }, + ) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_requests * per_call) + + +def _image_gen_token_usage(input_text, input_image, output_image, output_text): + return { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": { + "image_gen": { + "input_tokens": input_text + input_image, + "output_tokens": output_image + output_text, + "total_tokens": input_text + input_image + output_image + output_text, + "input_tokens_details": {"image_tokens": input_image, "text_tokens": input_text}, + "output_tokens_details": {"image_tokens": output_image, "text_tokens": output_text}, + } + }, + } + + +def _expected_image_gen_token_cost(model_info, input_text, input_image, output_image, output_text): + return ( + input_text * (model_info.get("input_cost_per_token") or 0) + + input_image * (model_info.get("input_cost_per_image_token") or 0) + + output_image * (model_info.get("output_cost_per_image_token") or 0) + + output_text * (model_info.get("output_cost_per_token") or 0) + ) + + +def test_responses_image_tool_model_from_tools_bills_token_usage(local_model_cost_map): + """The image tool's tools[].model is used and usage.tool_usage.image_gen tokens bill at that model's rates.""" + tool_model = "gpt-image-2" + model_info = litellm.get_model_info(tool_model, custom_llm_provider="openai") + tools = [{"type": "image_generation", "model": tool_model, "quality": "low", "size": "1024x1024"}] + output = [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "low", + "size": "1024x1024", + "result": "AAAA", + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ] + + response = _openai_responses_response("gpt-5", output, usage=_image_gen_token_usage(10, 0, 50, 0), tools=tools) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost > 0 + assert cost == pytest.approx(_expected_image_gen_token_cost(model_info, 10, 0, 50, 0)) + + response_more_tokens = _openai_responses_response( + "gpt-5", output, usage=_image_gen_token_usage(10, 0, 80, 0), tools=tools + ) + cost_more = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response_more_tokens, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost_more != cost + assert cost_more == pytest.approx(_expected_image_gen_token_cost(model_info, 10, 0, 80, 0)) + + +def test_responses_zero_image_gen_tokens_fall_back_to_per_image_pricing(local_model_cost_map): + """An all-zero image_gen usage block keeps the per-image path for the tool's model/quality/size.""" + tool_model = "gpt-image-1" + quality = "low" + size = "1024x1024" + tools = [{"type": "image_generation", "model": tool_model, "quality": quality, "size": size}] + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": quality, + "size": size, + "result": "AAAA", + } + ], + usage=_image_gen_token_usage(0, 0, 0, 0), + tools=tools, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + from litellm.cost_calculator import default_image_cost_calculator + + assert cost == pytest.approx( + default_image_cost_calculator( + model=tool_model, + custom_llm_provider="openai", + quality=quality, + n=1, + size=size, + ) + ) + assert cost > 0 + + +def test_responses_auto_size_image_generation_call_billed_at_default_size(local_model_cost_map): + """An image_generation_call with size "auto" bills at the default size instead of erroring to $0.""" + from litellm.cost_calculator import default_image_cost_calculator + + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "auto", + "result": "AAAA", + } + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx( + default_image_cost_calculator( + model="gpt-image-1", + custom_llm_provider="openai", + quality="high", + n=1, + size=None, + ) + ) + assert cost > 0 + + +def test_responses_image_gen_total_without_token_details_falls_back_to_per_image(local_model_cost_map): + """A positive image_gen total with no token details falls back to per-image pricing, not $0.""" + tool_model = "gpt-image-1" + quality = "low" + size = "1024x1024" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": quality, + "size": size, + "result": "AAAA", + } + ], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"image_gen": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}}, + }, + tools=[{"type": "image_generation", "model": tool_model, "quality": quality, "size": size}], + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + from litellm.cost_calculator import default_image_cost_calculator + + assert cost == pytest.approx( + default_image_cost_calculator( + model=tool_model, + custom_llm_provider="openai", + quality=quality, + n=1, + size=size, + ) + ) + assert cost > 0 @pytest.mark.parametrize( ("model_key", "web_search_requests", "medium_cost_per_query", "expected_cost"), [ From 79ca820933a3db556759c14e3ee0a01763a4deaf Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:19:55 +0000 Subject: [PATCH 06/82] fix(count_tokens): return a 400 naming Anthropic blocks the OpenAI and Gemini count can't convert (#46772) * fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does * test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends * fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types * fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them * fix(count_tokens): return a 400 naming Anthropic blocks the OpenAI and Gemini count can't convert * test(count_tokens): compare tools by value in the public api test * fix(count_tokens): reject unconvertible blocks nested inside a tool_result * fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field. * test(count_tokens): expect the image detail the adapter now sends after merging main * fix(count_tokens): fetch Gemini remote media without blocking the event loop * fix(count_tokens): type normalize_native_messages without the unimported Any On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError * chore(count_tokens): give the Gemini count helpers' list returns a reason * fix(count_tokens): normalize Anthropic messages for Gemini counts the way /v1/messages does * fix(count_tokens): count Gemini bodies after inference's history rewrites and answer converter errors with 400 * test(count_tokens): cover the Gemini count error paths and compaction with a text-block system --------- Co-authored-by: Mrinal Chanshetty --- .../pass_through/convertible_blocks.py | 106 ++++++++++ litellm/llms/gemini/common_utils.py | 4 + .../responses/count_tokens/token_counter.py | 6 +- litellm/llms/vertex_ai/common_utils.py | 4 + .../llms/vertex_ai/count_tokens/context.py | 9 + .../proxy/anthropic_endpoints/endpoints.py | 22 +- .../test_openai_count_tokens_tools_wire.py | 42 ++++ .../test_vertex_gemini_count_tokens_wire.py | 47 +++++ .../pass_through/test_convertible_blocks.py | 192 ++++++++++++++++++ tests/unit/proxy/test_proxy_token_counter.py | 37 ++-- 10 files changed, 430 insertions(+), 39 deletions(-) create mode 100644 litellm/llms/anthropic/pass_through/convertible_blocks.py create mode 100644 tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py diff --git a/litellm/llms/anthropic/pass_through/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py new file mode 100644 index 00000000000..543a6efd4c4 --- /dev/null +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -0,0 +1,106 @@ +from collections.abc import Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.types.utils import TokenCountMessagesFormat + +_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType( + { + "user": frozenset({"text", "image", "document", "tool_result"}), + "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}), + } +) +_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"}) +_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"}) + + +_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) + + +def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]: + try: + return _JSON_MESSAGES.validate_python(messages) + except ValidationError: + return () + + +def _source_label(block_type: str, source: JsonValue) -> str | None: + if not isinstance(source, dict): + return f"{block_type} without a source" + match source.get("type"): + case "base64" if source.get("data"): + return None + case "url" if source.get("url"): + return None + case "base64" | "url" as source_type: + return f"{block_type} with an empty {source_type} source" + case str() as source_type: + return f"{block_type} with a {source_type} source" + case _: + return f"{block_type} with an untyped source" + + +def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None: + if not isinstance(block, dict): + return None + block_type: Final = block.get("type") + if not isinstance(block_type, str): + return "a block without a type" + if block_type not in convertible: + return block_type + return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None + + +def _tool_result_labels(block: JsonValue) -> Iterator[str]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return iter(()) + inner: Final = block.get("content") + if not isinstance(inner, list): + return iter(()) + return ( + f"{label} inside a tool_result" + for item in inner + if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None + ) + + +def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]: + label: Final = _block_label(block, convertible) + return chain(() if label is None else (label,), _tool_result_labels(block)) + + +def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]: + role: Final = message.get("role") + convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None + content: Final = message.get("content") + if convertible is None or not isinstance(content, list): + return iter(()) + return chain.from_iterable(_content_labels(block, convertible) for block in content) + + +def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + normalized: Final = _json_messages(normalize_native_messages(messages)) + return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized)))) + + +def unconvertible_blocks_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> str | None: + match messages_format: + case "openai": + return None + case "anthropic": + labels: Final = unconvertible_anthropic_blocks(messages or ()) + return ( + "count_tokens can't convert these content blocks for this model, so the count would leave them out: " + + ", ".join(labels) + if labels + else None + ) + case _: + assert_never(messages_format) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index bac5824b981..3d04954dd41 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -507,6 +507,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): gemini_count_contents, inference_request_view, provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) @@ -514,6 +515,9 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index b97e1329675..43b3ae9f6f1 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -10,6 +10,7 @@ from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) @@ -91,14 +92,15 @@ class OpenAITokenCounter(BaseTokenCounter): if messages_format == "anthropic" else None ) - if malformed_error is not None: + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: return TokenCountResponse( total_tokens=0, request_model=request_model, model_used=model_to_use, tokenizer_type="openai_api", error=True, - error_message=malformed_error, + error_message=request_error, status_code=400, ) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 4c0026783f7..98bc679a78e 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1362,6 +1362,7 @@ class VertexAITokenCounter(BaseTokenCounter): gemini_count_contents, inference_request_view, provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter @@ -1370,6 +1371,9 @@ class VertexAITokenCounter(BaseTokenCounter): count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) resolved_contents: Final = ( contents diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 3f9289c8244..e094f4c5c0c 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -138,6 +138,15 @@ def vertex_gemini_count_context( ) +def unconvertible_messages_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> VertexGeminiCountContextError | None: + from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error + + error_message: Final = unconvertible_blocks_error(messages, messages_format) + return None if error_message is None else VertexGeminiCountContextError(error_message=error_message) + + def count_context_error_response( error: VertexGeminiCountContextError, request_model: str, model_used: str ) -> TokenCountResponse: diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1f17f0a42d2..4d0eb6a166a 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -43,7 +43,6 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.types.utils import TokenCountResponse router: Final = APIRouter() @@ -361,30 +360,15 @@ async def count_tokens( messages_format="anthropic", output_format=_requested_output_format(data), ) - _token_response_dict: dict = {} - if isinstance(token_response, TokenCountResponse): - _token_response_dict = token_response.model_dump() - elif isinstance(token_response, dict): - _token_response_dict = token_response - - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - except HTTPException: raise except ProxyException as e: - status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500 - detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=e.message, - ) - raise HTTPException( - status_code=status_code, - detail=detail, - ) + return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + else: + return {"input_tokens": token_response.total_tokens} @router.post( diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 3ab5826006f..2b23ccdf4f2 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -321,6 +321,48 @@ def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( assert count_body["model"] == inference_body["model"] +_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, + ], +} + + +def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY}) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert wire.drain() == () + + def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_count_response) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 7fef79845a1..5fe82d37e9e 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -254,6 +254,53 @@ def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gatew assert _counted_contents(count_body) == generated_body["contents"] +_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, +] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert peer.drain() == () + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py new file mode 100644 index 00000000000..b5301b413ff --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py @@ -0,0 +1,192 @@ +import asyncio +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.pass_through.convertible_blocks import ( + unconvertible_anthropic_blocks, + unconvertible_blocks_error, +) +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents +from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues + +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + +_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} +_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + +_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ("user", {"type": "text", "text": "hi"}), + ("user", {"type": "image", "source": _PNG}), + ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}), + ("user", {"type": "document", "source": _PDF}), + ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}), + ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}), + ("assistant", {"type": "text", "text": "hi"}), + ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}), + ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}), +) + +_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = ( + ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"), + ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"), + ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"), + ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"), + ("user", {"type": "document"}, "document without a source"), + ("user", {"text": "untyped"}, "a block without a type"), + ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"), + ( + "assistant", + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}}, + "server_tool_use", + ), +) + + +def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose + + +def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]: + messages: Final = _adapter_messages(role, block) + chat: Final = _JSON_OBJECTS.validate_python( + LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages) + ) + responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages) + chat_kept: Final = any( + message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat + ) + return chat_kept, bool(responses) + + +@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS) +def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None: + assert _converted_by_both_adapters(role, block) == (True, True) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == () + + +@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS) +def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None: + assert _converted_by_both_adapters(role, block) == (False, False) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,) + + +def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None: + messages: Final = [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]}, + {"role": "assistant", "content": [{"type": "search_result"}]}, + ] + assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result") + + +def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None: + messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def test_messages_that_are_not_json_are_left_to_the_count() -> None: + messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]}, + ) + + +def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]: + messages: Final = _tool_turn(inner) + gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {})) + openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input( + [dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose + ) + return json.dumps(gemini, default=str), json.dumps(openai, default=str) + + +_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = ( + ({"type": "text", "text": "forty-two"}, "forty-two"), + ({"type": "image", "source": _PNG}, _PNG["data"]), + ({"type": "document", "source": _PDF}, _PDF["data"]), +) + +_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = ( + ( + { + "type": "search_result", + "source": "https://example.com", + "title": "t", + "content": [{"type": "text", "text": "search-snippet"}], + }, + "search-snippet", + "search_result inside a tool_result", + ), + ( + {"type": "document", "source": {"type": "text", "data": "text-notes"}}, + "text-notes", + "document with a text source inside a tool_result", + ), + ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"), +) + + +@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker in gemini + assert marker in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == () + + +@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker not in gemini + assert marker not in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,) + + +def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": encrypted_content, + "snippet": "LiteLLM ships a weekly release candidate.", + } + ], + }, + ], + }, + ) + + +def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None: + assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None + + +def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None: + error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic") + assert error is not None + assert error.endswith("server_tool_use, web_search_tool_result") diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 8cd6887d891..9bed719227f 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,6 +2,8 @@ # 1. Generate a Key, and use it to make a call +import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -14,7 +16,8 @@ load_dotenv() # this file is to test litellm/proxy -from fastapi import HTTPException, Request +from fastapi import Request +from fastapi.responses import JSONResponse import litellm from litellm import Router @@ -1035,6 +1038,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1062,14 +1066,11 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - # Verify HTTP status code is correct - assert exc_info.value.status_code == 400 - - # Verify error is in Anthropic format - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "invalid_request_error" assert detail["error"]["message"] == "Input is too long for requested model." @@ -1087,6 +1088,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1114,12 +1116,11 @@ async def test_anthropic_endpoint_403_permission_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 403 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "permission_error" assert detail["error"]["message"] == "Bearer Token has expired" @@ -1137,6 +1138,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1164,12 +1166,11 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 429 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "rate_limit_error" assert detail["error"]["message"] == "Rate limit exceeded" From e74e4e569331c1d99757ae6118685d8826c6f7c6 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:31:10 -0700 Subject: [PATCH 07/82] fix(vertex-ai): drop cached flex price from gemini 3 preview rows (#46826) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 138 ++++++++---------- model_prices_and_context_window.json | 138 ++++++++---------- 2 files changed, 118 insertions(+), 158 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a204c49967e..b744032fb6c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29256,22 +29256,32 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29289,33 +29299,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29428,15 +29423,25 @@ "vertex_ai/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29454,6 +29459,7 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -29463,23 +29469,7 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -29807,22 +29797,32 @@ "output_cost_per_video_token": 1e-06 }, "vertex_ai/gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29840,33 +29840,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -31666,9 +31651,15 @@ "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -31676,6 +31667,10 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31693,6 +31688,7 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -31703,23 +31699,7 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a204c49967e..b744032fb6c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29256,22 +29256,32 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29289,33 +29299,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29428,15 +29423,25 @@ "vertex_ai/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29454,6 +29459,7 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -29463,23 +29469,7 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -29807,22 +29797,32 @@ "output_cost_per_video_token": 1e-06 }, "vertex_ai/gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29840,33 +29840,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -31666,9 +31651,15 @@ "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -31676,6 +31667,10 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31693,6 +31688,7 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -31703,23 +31699,7 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, From 5570d4268ce9eee71230ef984a1fa02be2479acb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:34:48 -0700 Subject: [PATCH 08/82] fix(proxy): only release the legacy limiter's slot from an admitted Responses WebSocket (#46822) The shared disconnect release dispatched to the legacy limiter, which decrements unconditionally, so a rejected realtime attempt freed a slot it never took. The generic release is v3-only again and the Responses WebSocket close releases through its own method Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/response_api_endpoints/endpoints.py | 2 +- litellm/proxy/utils.py | 10 +++++++- .../hooks/test_parallel_request_limiter.py | 25 ++++++++++++++++++- .../response_api_endpoints/test_endpoints.py | 4 +-- 4 files changed, 36 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index b7f57d93a9b..a44f1d10fda 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1906,7 +1906,7 @@ async def _responses_websocket_session( ) await websocket.close(code=1011, reason="Internal server error") finally: - await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5d2ebe73cac..be6532ba36b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4405,12 +4405,20 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests_on_disconnect + async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Release the slot a Responses WebSocket took at admission. Only call this once that admission succeeded: + the legacy limiter keeps no record of its acquisitions and decrements unconditionally.""" + limiter: Final = self.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + return + await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + async def enforce_mcp_server_rate_limits( self, user_api_key_dict: UserAPIKeyAuth | None, diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index f17cad56751..08840978169 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -15,7 +15,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) -from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) @@ -353,3 +353,26 @@ async def test_websocket_turn_callbacks_keep_the_sockets_slot_until_the_disconne await handler.async_release_max_parallel_requests_on_disconnect(UserAPIKeyAuth(api_key=api_key)) after_disconnect = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) assert after_disconnect["current_requests"] == 0, after_disconnect + + +@pytest.mark.asyncio +async def test_a_rejected_realtime_attempt_does_not_free_a_slot_it_never_took(): + api_key = hash_token("sk-at-limit") + handler = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock) + bucket = f"{api_key}::{FROZEN_INSTANT.strftime('%Y-%m-%d-%H-%M')}::request_count" + await handler.internal_usage_cache.async_set_cache( + key=bucket, + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler + auth = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth) + after_rejection = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_rejection["current_requests"] == 1, after_rejection + + await proxy_logging_obj.arelease_responses_websocket_slot(auth) + after_socket_close = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_socket_close["current_requests"] == 0, after_socket_close diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0eaf619ba70..b047c29243a 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1952,7 +1952,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -2017,7 +2017,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( From cb7e6a7f4c90f196e5cedbf5ee0dd16a71c814ef Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 20:35:35 -0700 Subject: [PATCH 09/82] test(e2e): mixed_size_requests_2rps perf case (coding-agent bodies 5k to 1m tokens with tools, JWT and key callers, two guardrails) (#46717) * test(e2e): mixed_size_requests_2rps perf case (coding-agent bodies, JWT and key callers, two guardrails) * test(e2e_harness): test_observe builds DbCounters with statements and expects the background Postgres line --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../perf/cases/mixed_size_requests_2rps.py | 74 +++++++++++++++++++ tests/e2e_harness/perf/setup/test_observe.py | 11 ++- 2 files changed, 82 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/perf/cases/mixed_size_requests_2rps.py diff --git a/tests/e2e/perf/cases/mixed_size_requests_2rps.py b/tests/e2e/perf/cases/mixed_size_requests_2rps.py new file mode 100644 index 00000000000..41b61cdc987 --- /dev/null +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -0,0 +1,74 @@ +"""Mixed-size requests at 2 rps on the small gateway: three real model groups over the mock providers, +two teams plus users in no team, JWT and key callers, a pre and +a post guardrail, half the requests streaming, short chats, long contexts and tool calls. Budgets +are real money at cost-map prices (a 335k-token `claude-sonnet-5` request is about $0.70 and the +run spends roughly $600), so they are sized so no caller runs dry mid-run.""" + +from __future__ import annotations + +from typing import Final + +from setup.performance_test_case import ( + LONG, + SHORT, + Guardrail, + Limit, + MockLatency, + PerformanceTestCase, + Team, + Traffic, + User, + VirtualKey, +) + +CASE: Final = PerformanceTestCase( + name="mixed_size_requests_2rps", + teams=( + Team("team-a", models=("gpt-6-luna", "claude-sonnet-5"), max_budget=5000.0, rpm_limit=6000), + Team("team-b", models=("claude-sonnet-5-bedrock", "claude-sonnet-5"), max_budget=5000.0), + ), + users=( + User("alice", team="team-a"), + User("bob", team="team-a"), + User("carol", team="team-a"), + User("dave", team="team-b"), + User("erin", team="team-b"), + User("frank", team="team-b"), + User("grace"), + User("heidi"), + User("ivan"), + User("judy", max_budget=1000.0), + ), + virtual_keys=( + VirtualKey("key-a1", team="team-a", models=("gpt-6-luna",), rpm_limit=1200), + VirtualKey("key-a2", team="team-a", user="alice"), + VirtualKey("key-b1", team="team-b"), + VirtualKey("key-b2", team="team-b", user="dave"), + VirtualKey("key-plain"), + ), + model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"), + mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25), + guardrails=(Guardrail("perf-pre", mode="pre_call"), Guardrail("perf-post", mode="post_call")), + traffic=Traffic( + seed=7, + rps=2.0, + total_requests=1200, + endpoints={"chat_completions": 0.5, "responses": 0.25, "messages": 0.25}, + shapes={SHORT: 0.75, LONG: 0.25}, + stream_share=0.5, + intervene_share=0.0, + block_share=0.0, + concurrency=100, + ), + limits=( + Limit("unexpected_status_count", max=3), + Limit("gateway_added_seconds", percentile=50, max=0.58), + Limit("gateway_added_seconds", percentile=99, max=6.1), + Limit("time_to_first_token_seconds", percentile=50, max=2.7), + Limit("time_to_first_token_seconds", percentile=99, max=8.8), + Limit("memory_growth_megabytes", max=600), + Limit("cpu_seconds_per_request", max=1.26), + Limit("postgres_calls_per_request", max=3.5), + Limit("redis_calls_per_request", max=21.0), + ), +) diff --git a/tests/e2e_harness/perf/setup/test_observe.py b/tests/e2e_harness/perf/setup/test_observe.py index 998d7d9b87e..f38e7927fb7 100644 --- a/tests/e2e_harness/perf/setup/test_observe.py +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -5,7 +5,7 @@ from typing import Final from cases.small_bodies_5rps import CASE from memory_readings import RssCapture -from setup.datastore import DbCounters +from setup.datastore import DbCounters, StatementCalls from setup.observe import Observation, evaluate, junit_properties from setup.performance_test_case import Limit from setup.wire import WireRecord @@ -35,6 +35,10 @@ def served(seq: int, total_s: float) -> WireRecord: ) +def _key_lookup(calls: int) -> StatementCalls: + return StatementCalls(queryid=1, query='SELECT * FROM "LiteLLM_VerificationToken" WHERE token = $1', calls=calls) + + def test_junit_properties_carry_each_observed_metric_and_each_limit_including_unobserved_ones() -> None: case: Final = replace( CASE, @@ -50,8 +54,8 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un rss_before=NO_RSS, rss_after=NO_RSS, wall_seconds=3.0, - db_before=DbCounters(commits=100), - db_after=DbCounters(commits=130), + db_before=DbCounters(commits=100, statements=(_key_lookup(calls=100),)), + db_after=DbCounters(commits=130, statements=(_key_lookup(calls=130),)), ) rows, _ = evaluate(case, observation) @@ -64,5 +68,6 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un "perf_gateway_added_seconds_p50_limit": 0.625, "perf_gateway_added_seconds_p99": 0.75, "perf_postgres_calls_per_request": 10.0, + "perf_postgres_background_calls_per_request": 0.0, "perf_redis_calls_per_request_limit": 35.0, } From 0620f1a51586b5d0b0de2ab7653dc76de935ea0d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 15 Oct 2026 21:44:58 -0700 Subject: [PATCH 10/82] fix(ui): validate the Add Model form when Test Connect is clicked (#46834) * fix(ui): gate Add Model's Test Connect on the form's own validation Test Connect stays disabled until the Add Model form would pass validation, with a tooltip saying why. Provider password, text and textarea inputs now render empty after a successful add resets the form, and an empty prepared request shows a plain message instead of a destructuring crash * fix(ui): validate the Add Model form when Test Connect is clicked Test Connect stays enabled and, on click, runs the same validation Add Model runs. An incomplete form now shows the inline required errors instead of opening a crashing connection test, and a valid form tests as before. This replaces the disabled-button gate, which read a validity flag that went stale after programmatic form updates * test(ui): cover a blank required Azure API Base on Test Connect --- .../AddModelForm.integration.test.tsx | 225 +++++++++++++++++- .../src/components/add_model/AddModelForm.tsx | 13 +- .../add_model/model_connection_test.test.tsx | 13 + .../add_model/model_connection_test.tsx | 8 +- .../provider_specific_fields.test.tsx | 59 +++++ .../add_model/provider_specific_fields.tsx | 9 +- 6 files changed, 310 insertions(+), 17 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 681217cc918..cdc681a1c91 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -8,9 +8,9 @@ import { renderWithProviders, } from "../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import type { Team } from "../key_team_helpers/key_list"; -import { credentialCreateCall, modelCreateCall, type CredentialItem } from "../networking"; +import { credentialCreateCall, modelCreateCall, testConnectionRequest, type CredentialItem } from "../networking"; import { Providers } from "../provider_info_helpers"; import { projectMountedValues, useMountRegistry, type MountedFormValues } from "../common_components/MountedFormField"; import { useForm } from "react-hook-form"; @@ -53,7 +53,17 @@ vi.mock("../networking", async () => { provider_display_name: "OpenAI", litellm_provider: "openai", default_model_placeholder: "gpt-3.5-turbo", - credential_fields: [], + credential_fields: [{ key: "api_key", label: "OpenAI API Key", field_type: "password", required: true }], + }, + { + provider: "Azure", + provider_display_name: "Azure", + litellm_provider: "azure", + default_model_placeholder: "my-deployment", + credential_fields: [ + { key: "api_base", label: "API Base", field_type: "text", required: true }, + { key: "api_key", label: "Azure API Key", field_type: "password" }, + ], }, { provider: "MICROSOFT_365_COPILOT", @@ -98,7 +108,17 @@ vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ provider_display_name: "OpenAI", litellm_provider: "openai", default_model_placeholder: "gpt-3.5-turbo", - credential_fields: [], + credential_fields: [{ key: "api_key", label: "OpenAI API Key", field_type: "password", required: true }], + }, + { + provider: "Azure", + provider_display_name: "Azure", + litellm_provider: "azure", + default_model_placeholder: "my-deployment", + credential_fields: [ + { key: "api_base", label: "API Base", field_type: "text", required: true }, + { key: "api_key", label: "Azure API Key", field_type: "password" }, + ], }, { provider: "MICROSOFT_365_COPILOT", @@ -252,6 +272,22 @@ const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmi }; }; +type User = ReturnType; + +const pickProvider = async (user: User, name: RegExp) => { + await user.click(screen.getByRole("combobox", { name: /provider/i })); + await user.click(await screen.findByRole("option", { name })); +}; + +const pickModel = async (user: User, name: string) => { + await user.click(await screen.findByPlaceholderText("Select models")); + await user.click(await screen.findByRole("option", { name })); + await user.keyboard("{Escape}"); +}; + +const fillApiKey = (label: string) => + fireEvent.change(screen.getByLabelText(new RegExp(label)), { target: { value: "sk-test" } }); + describe("AddModelForm", () => { it("should render", async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); @@ -444,10 +480,15 @@ describe("AddModelForm", () => { it("shows only the Close button in the connection test dialog footer", async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); renderWithProviders(); + await screen.findByText("Provider"); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); - await userEvent.click(await screen.findByTestId("test-connect-btn")); + await user.click(screen.getByTestId("test-connect-btn")); const dialog = await screen.findByRole("dialog"); const footer = dialog.querySelector('[data-slot="dialog-footer"]'); @@ -786,6 +827,180 @@ describe("AddModelForm", () => { }); }); + describe("Test Connect validates the form the way Add Model does", () => { + beforeEach(() => { + vi.mocked(testConnectionRequest).mockClear(); + }); + + const renderForm = async (selectedProvider: Providers = Providers.OpenAI) => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + const props = { ...createTestProps(), selectedProvider }; + renderWithProviders(); + await screen.findByText("Provider"); + return { user, props, clickTestConnect: () => user.click(screen.getByTestId("test-connect-btn")) }; + }; + + const expectTestRun = async (litellmParams: Record) => { + expect(await screen.findByTestId("connection-success-msg")).toBeInTheDocument(); + expect(testConnectionRequest).toHaveBeenCalledWith( + "test-access-token", + expect.objectContaining(litellmParams), + expect.anything(), + undefined, + ); + }; + + const expectBlockedWith = async (inlineError: string) => { + expect(await screen.findByText(inlineError)).toBeInTheDocument(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + expect(testConnectionRequest).not.toHaveBeenCalled(); + }; + + it("shows the inline required errors instead of testing a fresh form", async () => { + const { clickTestConnect } = await renderForm(); + + await clickTestConnect(); + + await expectBlockedWith("Please enter at least one model."); + expect(screen.getAllByText("Required")).toHaveLength(2); + }); + + it("runs the connection test once provider, auth and model are filled in", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + + await clickTestConnect(); + + await expectTestRun({ api_key: "sk-test", model: "gpt-4" }); + }); + + it("shows emptied inputs and inline errors instead of testing after a successful Add Model", async () => { + const { user, props, clickTestConnect } = await renderForm(); + props.handleOk.mockImplementation(async () => { + props.form.reset({ litellm_credential_name: null }); + return true; + }); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + await user.click(screen.getByTestId("add-model-btn")); + await waitFor(() => expect(screen.getByLabelText(/OpenAI API Key/)).toHaveValue("")); + expect(props.handleOk).toHaveBeenCalledOnce(); + + await clickTestConnect(); + + await expectBlockedWith("Please enter at least one model."); + }); + + it("runs the connection test for the provider wildcard", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + await pickModel(user, "All OpenAI Models (Wildcard)"); + fillApiKey("OpenAI API Key"); + + await clickTestConnect(); + + await expectTestRun({ model: "openai/*" }); + }); + + it("asks for the custom model name before testing a custom model", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "Custom Model Name (Enter below)"); + + await clickTestConnect(); + await expectBlockedWith("Please enter a custom model name."); + + fireEvent.change(screen.getByPlaceholderText("Enter custom model name"), { target: { value: "my-fine-tune" } }); + await clickTestConnect(); + + await expectTestRun({ model: "my-fine-tune" }); + }); + + it("runs the connection test for an Azure deployment name", async () => { + const { user, clickTestConnect } = await renderForm(Providers.Azure); + await pickProvider(user, /Azure/); + fireEvent.change(screen.getByLabelText(/LiteLLM Model Name/), { target: { value: "my-deployment" } }); + + await clickTestConnect(); + await expectBlockedWith("Required"); + + fireEvent.change(screen.getByLabelText(/API Base/), { target: { value: "https://example.openai.azure.com" } }); + await clickTestConnect(); + + await expectTestRun({ model: "azure/my-deployment", api_base: "https://example.openai.azure.com" }); + }); + + it("asks for a credential before testing a credential-only auth type", async () => { + const errorToast = vi.spyOn(toast, "error"); + const { user, clickTestConnect } = await renderForm(Providers.MICROSOFT_365_COPILOT); + await pickProvider(user, /Microsoft 365 Copilot/); + await pickModel(user, "gpt-4"); + expect(await screen.findByRole("button", { name: "Create credential" })).toBeInTheDocument(); + + await clickTestConnect(); + expect(errorToast).toHaveBeenCalledWith("Create or select a credential for this auth type"); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + expect(testConnectionRequest).not.toHaveBeenCalled(); + errorToast.mockRestore(); + + await user.click(screen.getByPlaceholderText("Select or search for existing credentials")); + await user.click(await screen.findByRole("option", { name: /test-credential/ })); + await clickTestConnect(); + + await expectTestRun({ litellm_credential_name: "test-credential" }); + }); + + it("runs the connection test with a credential just created from the credential dialog", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + await pickModel(user, "gpt-4"); + await user.click(screen.getByRole("button", { name: "Use workload identity federation" })); + fireEvent.change(screen.getByLabelText("Credential Name:"), { target: { value: "openai-federated" } }); + fireEvent.change(screen.getByLabelText(/Service Account ID/), { target: { value: "svc_new" } }); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + await waitFor(() => expect(screen.queryByLabelText(/OpenAI API Key/)).not.toBeInTheDocument()); + + await clickTestConnect(); + + await expectTestRun({ litellm_credential_name: "openai-federated", model: "gpt-4" }); + }); + + it("runs the connection test after the public model name is edited in the mapping table", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + const publicName = screen.getByDisplayValue("gpt-4"); + await user.clear(publicName); + const apiKey = screen.getByLabelText(/OpenAI API Key/); + await user.clear(apiKey); + await user.type(apiKey, "sk-rotated"); + await user.type(publicName, "my-gpt"); + + await clickTestConnect(); + + await expectTestRun({ api_key: "sk-rotated", model: "gpt-4" }); + }); + + it("shows the mapping error instead of testing when the public model name is cleared", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + fireEvent.change(screen.getByDisplayValue("gpt-4"), { target: { value: "" } }); + + await clickTestConnect(); + + await expectBlockedWith("All model mappings must have valid public names"); + }); + }); + describe("cache control bindings reach the parent form store", () => { const renderWithForm = async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 5313df6153a..7e90991523c 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -134,15 +134,20 @@ const AddModelForm: React.FC = ({ setIsCredentialModalOpen(true); }; - const handleSubmit = async () => { - if (selectedAuthType?.credentialOnly && !selectedCredentialName) { + const rejectMissingCredential = (): boolean => { + const missing = selectedAuthType?.credentialOnly === true && !selectedCredentialName; + if (missing) { toast.error("Create or select a credential for this auth type"); - return false; } - return handleOk(); + return missing; }; + const handleSubmit = async () => (rejectMissingCredential() ? false : handleOk()); + const handleTestConnection = async () => { + if (rejectMissingCredential() || !(await form.trigger(registry.mountedNames() as string[]))) { + return; + } setIsTestingConnection(true); setConnectionTestId(`test-${Date.now()}`); setIsResultModalVisible(true); diff --git a/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx b/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx index c436a24474c..3bc867f843c 100644 --- a/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx @@ -93,6 +93,19 @@ describe("ModelConnectionTest", () => { expect(toast.success).toHaveBeenCalledWith("Copied to clipboard"); }); + it("asks for the required fields instead of crashing when nothing could be prepared", async () => { + vi.mocked(prepareModelAddRequest).mockResolvedValue([] as never); + + render(); + await finishConnectionTest(); + + expect(screen.getByTestId("connection-failure-msg")).toHaveTextContent("Connection to this model failed"); + expect( + screen.getByText("Required fields are not filled in. Fill them in and test the connection again."), + ).toBeInTheDocument(); + expect(testConnectionRequest).not.toHaveBeenCalled(); + }); + it("shows a preparation failure without sending a connection request", async () => { vi.mocked(prepareModelAddRequest).mockResolvedValue(null as never); diff --git a/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx b/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx index f6ef86a77ca..6fdc0461482 100644 --- a/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx @@ -43,8 +43,12 @@ const ModelConnectionTest: React.FC = ({ try { const result = await prepareModelAddRequest(formValues, accessToken, null); - if (!result) { - setError("Failed to prepare model data. Please check your form inputs."); + if (!result?.length) { + setError( + result + ? "Required fields are not filled in. Fill them in and test the connection again." + : "Failed to prepare model data. Please check your form inputs.", + ); setIsSuccess(false); setIsLoading(false); return; diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx index 09d48f828e1..3a58da9248a 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx @@ -112,6 +112,13 @@ vi.mock("../networking", async () => { }, ], }, + { + provider: "Oracle", + provider_display_name: Providers.Oracle, + litellm_provider: "oci", + default_model_placeholder: "oci/xai.grok-4", + credential_fields: [{ key: "oci_key", label: "OCI Private Key (PEM)", field_type: "textarea", required: true }], + }, { provider: "MICROSOFT_365_COPILOT", provider_display_name: Providers.MICROSOFT_365_COPILOT, @@ -182,6 +189,15 @@ const VertexCredentialsProbe = () => { return {String(watch("vertex_credentials") ?? "")}; }; +const ResetFormButton = () => { + const form = useFormContext(); + return ( + + ); +}; + const ValidationForm = ({ children }: { readonly children: ReactNode }) => { const form = useFormContext(); return ( @@ -642,6 +658,49 @@ describe("ProviderSpecificFields", () => { expect(await screen.findAllByText("Required")).toHaveLength(3); }); + it("renders typed password and text inputs empty after the form is reset", async () => { + render( + + + + + + , + ); + const apiKey = await screen.findByLabelText("OpenAI API Key"); + const organization = screen.getByLabelText("OpenAI Organization ID"); + const apiBase = screen.getByLabelText(/API Base/); + expect(apiBase).toHaveValue("https://api.openai.com/v1"); + fireEvent.change(apiKey, { target: { value: "sk-typed" } }); + fireEvent.change(organization, { target: { value: "org-typed" } }); + fireEvent.change(apiBase, { target: { value: "https://proxy.example.test/v1" } }); + expect(apiKey).toHaveValue("sk-typed"); + + fireEvent.click(screen.getByRole("button", { name: "Reset form" })); + + await waitFor(() => expect(apiKey).toHaveValue("")); + expect(organization).toHaveValue(""); + expect(apiBase).toHaveValue("https://api.openai.com/v1"); + }); + + it("renders a typed textarea input empty after the form is reset", async () => { + render( + + + + + + , + ); + const privateKey = await screen.findByLabelText(/OCI Private Key/); + fireEvent.change(privateKey, { target: { value: "-----BEGIN PRIVATE KEY-----" } }); + expect(privateKey).toHaveValue("-----BEGIN PRIVATE KEY-----"); + + fireEvent.click(screen.getByRole("button", { name: "Reset form" })); + + await waitFor(() => expect(privateKey).toHaveValue("")); + }); + it("does not render an Auth Type selector for OpenAI", async () => { const queryClient = createQueryClient(); render( diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index a38b077270f..31a7cb02d6e 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -301,11 +301,10 @@ const ProviderSpecificFieldsContent: React.FC = ({ return (