mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
chore: litellm oss staging160626 (#30527)
* feat(ui): gate "Default Credentials" hint on /ui/login behind env flag (#30234) Adds LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT (and an equivalent general_settings.hide_default_credentials_hint) that suppresses the "By default, Username is admin and Password is your set LiteLLM Proxy MASTER_KEY" info card rendered on /ui/login and /fallback/login. Motivation: in production deployments operators set UI_USERNAME / UI_PASSWORD (or SSO), and the hardcoded hint becomes factually incorrect and is flagged by security scanners (Tenable WAS plugin 114625) as information disclosure. There is currently no way to suppress it without forking the dashboard. Behaviour: - Default is unchanged (hint shown), so existing deployments are unaffected. - New field hide_default_credentials_hint on the well-known UI config endpoint, populated from the env var or general_settings. - LoginPage.tsx conditionally renders the Alert based on the flag. Refs: BerriAI/litellm#30232 * fix(router): clean pattern_router state on upsert/delete (#29601) * fix(router): clean pattern_router state on upsert/delete PatternMatchRouter.add_pattern was append-only, and neither Router.upsert_deployment nor Router.delete_deployment removed the existing entry. Rotated-out api_keys stayed in the routing rotation for wildcard deployments (model_name with `*`) until proxy restart, silently defeating key rotation as an admin operation. The same leak applied to provider_default_deployment_ids and per-team pattern routers, and the patterns list grew unboundedly on every edit * test(router): direct unit tests for _remove_deployment_from_wildcard_state router_code_coverage.py greps test files for AST Call nodes and flagged the helper as untested because the existing coverage only exercised it transitively through upsert/delete. Adds two direct tests that pin the helper's contract (cleans across global pattern router, per-team routers with empty-router pop, and provider_default_deployment_ids; noop on falsy model_id) * fix(router): address Greptile review on pattern_router cleanup Widen PatternMatchRouter.remove_deployment annotation to Optional[str]; the implementation already handles None via the falsy guard and the unit test exercises it directly. Move _remove_deployment_from_wildcard_state up one level in upsert_deployment so it runs whenever the prior deployment is on the router, not only when the model_id is present in the fast-mapping index. The scenario is currently unreachable (get_deployment shares the same index), but the cleanup is idempotent so this is defensive against any future divergence between those code paths. * fix(router): widen _remove_deployment_from_wildcard_state to Optional[str] Moving the call out of the inner `deployment_id in deployment_fast_mapping` block in the previous commit lost mypy's narrowing of `deployment_id` from Optional[str] to str, tripping the lint CI. The helper already handles None via its falsy guard, so widening the annotation matches the actual contract. * fix(router): make delete_deployment wildcard cleanup symmetric with upsert After the previous commit moved _remove_deployment_from_wildcard_state out of the inner index-map guard in upsert_deployment, delete_deployment was still calling it only inside `if deployment_idx is not None`. Greptile flagged the asymmetry: under a desynced index_map, delete would silently leave the stale wildcard credential in pattern_router. Moves the cleanup call to the top of the try block, mirroring the upsert path. Cleanup is idempotent so the change is a no-op on the happy path. Adds a regression test that simulates the desync by removing the entry from model_id_to_deployment_index_map and asserts delete still clears pattern_router. * fix(pricing): add 1h cache-write cost for Anthropic Sonnet 4.5/4.6 (#30474) The native anthropic claude-sonnet-4-5/4-6 price-map entries were missing cache_creation_input_token_cost_above_1hr (and the >200K long-context sub-tier for 4.5), so 1-hour-TTL cache writes were costed at the 5-minute rate. Adds 6e-06 regular (and 1.2e-05 long-context) = 2x base input, matching the vertex_ai/azure_ai/bedrock siblings and the older claude-sonnet-4-20250514 entry. Adds a regression test. * fix(proxy): cancel upstream gemini request and release httpx connection on client disconnect (#30075) * fix(proxy): cancel upstream gemini request and release httpx connection on client disconnect - add _check_request_disconnection to common_request_processing; wrap llm_call as asyncio.Task so it can be cancelled; catch CancelledError and raise HTTPException(499) when client disconnects before LLM responds (non-streaming path) - pass raw httpx.Response into ModelResponseIterator in make_call/make_sync_call so the iterator holds a reference to the underlying connection - implement ModelResponseIterator.aclose() and .close(): close the line iterator then explicitly call response.aclose()/response.close() to release the httpx connection when the client drops mid-stream; errors are debug-logged, not raised - add tests for _check_request_disconnection (cancels task, graceful on exception, does not cancel when client stays connected) and base_process_llm_request 499 behavior; add TestModelResponseIteratorCleanup verifying aclose/close propagation through CustomStreamWrapper * fix(proxy): record 499 on streaming disconnect and cancel orphaned gather tasks Wire streaming generator cleanup to log client_disconnected with error_code 499 in spend logs, cancel pending during_call_hook tasks when the LLM call is cancelled on disconnect, and align the 600s poll limit comment with proxy_server. * fix: extract client disconnect logging helper to satisfy PLR0915 * fix: resolve mypy and code-quality CI failures for client disconnect logging Cast client disconnect error_information for mypy, only await pending gather tasks to avoid masking LLM errors, and add tests for the new logging helper and gather cleanup. * fix(proxy): harden gather cleanup so finally cannot mask LLM errors * fix(proxy): shield streaming disconnect logging and strip spoofable metadata Move streaming disconnect recording into a shielded cancel scope, add gather cleanup regression coverage for guardrail-converted cancels, and strip client_disconnected/error_information from user metadata at the proxy boundary. * fix(proxy): only map CancelledError to 499 for client disconnect Track when the disconnect poller cancels the LLM task and re-raise other CancelledError paths so graceful shutdown is not reported as HTTP 499. * fix(proxy): remove dead _check_request_disconnection helper Non-streaming client disconnect is handled by staging's cancel_on_disconnect path via _await_llm_call_cancelling_on_disconnect. Drop the unused is_disconnected poller and its unit tests; rename the remaining integration tests to TestDisconnectGatherCleanup. * feat(mistral): add mistral-medium-3-5 to model_prices_and_context_wind.. (#29303) * feat(mistral): add mistral-medium-3-5 to model_prices_and_context_window.json Mistral's docs page lists mistral-medium-3-5 as a new model offering. Pricing/specs sourced from Mistral's published model metadata: - input: $1.50 / 1M tokens - output: $7.50 / 1M tokens - context: 262,144 tokens - capabilities: vision, function calling, structured outputs, assistant prefill Adds entry: `mistral/mistral-medium-3-5`, mirroring the pattern used for the rest of the Mistral family. test(mistral): add model_info test for mistral-medium-3-5 + sync backup cost map - Mirror mistral/mistral-medium-3-5 entries into litellm/model_prices_and_context_window_backup.json so the bundled model cost map matches the canonical model_prices_and_context_window.json. - Add tests/test_litellm/test_mistral_medium_3_5_model_metadata.py covering pricing tiers, capability flags, context window, provider routing, and parity between the main and backup cost maps. - Point 'source' at the live Mistral models documentation page. * fix(ui): three small UI fixes — Gemini api_base + credential form reset + Mode badge (#30419) * fix(ui): three small UI fixes — Gemini api_base field + credential form reset + Mode badge Three independent fixes; bundled because they all touch the credential-form / logging-callbacks area. 1. expose api_base field on Google AI Studio credential form The runtime gemini provider supports custom api_base via `vertex_llm_base._check_custom_proxy`; the UI just needs to expose the field. Adds api_base to the Google_AI_Studio credential form ordered before api_key (matching OpenAI/Anthropic conventions). Default value matches the canonical Google AI Studio endpoint that LiteLLM's gemini provider talks to when api_base is unset, so leaving the default in the form behaves identically to leaving it blank. 2. reset credential form state when switching providers Switching the Provider select in AddCredentialModal / EditCredentialModal left the previous provider's field values populated. The form then submitted a mixed payload (e.g. Azure deployment fields under an OpenAI credential), producing confusing failures. Extract `getProviderFieldDefaults` helper and reset the form to it on provider change. Unit-tested via the extracted helper because Antd Select's portal/dropdown behaviour is unreliable in jsdom. 3. logging callbacks table reads backend `type` for Mode badge (#35) The `/get_callbacks` proxy endpoint returns each callback as `{name, type, variables}` where `type` is `"success"` or `"failure"`. The same callback name can appear twice (one per event class) and the two entries fire on disjoint events. `LoggingCallbacksTable` ignored `type` and read `record.mode` (always undefined), so every row fell back to the "Success" badge. A `generic_api` callback registered for both classes showed up as two identical "Success" rows + React duplicate-key warning. Read `record.type` first (fall back to `record.mode` for newly- added not-yet-server-acknowledged rows). Composite rowKey `${name}-${type ?? mode ?? 'success'}`. Removed leftover debug `console.log`. * fix(ui): drop api_base default_value to preserve Gemini v1alpha auto-routing Greptile P2 (PR #30419, threads on lines 1255-1256 of provider_create_fields.json): the api_base field's `default_value` was hard-coded to "https://generativelanguage.googleapis.com/v1beta". This: 1. Bakes v1beta into every credential record saved through the form, even when the user never touched the field. If LiteLLM's internal gemini default URL ever changes, those persisted credentials keep hitting the stale path. 2. Bypasses `_get_gemini_url`'s automatic version routing for Gemini 3+ models. That helper picks v1alpha for Gemini 3+ and v1beta for older models when api_base is unset. With the default pre-filled (and `_check_custom_proxy` then taking over because api_base is non-empty), Gemini 3+ requests get pinned to v1beta and may fail or behave unexpectedly — purely because the user accepted the visible default. Fix: set `default_value` to `null` and move the canonical URL guidance into the `placeholder` (visible to the user, never persisted) and an expanded tooltip. UX is unchanged — the URL is still shown in the greyed-out input — but the auto-version-routing path stays default. Updated test_google_ai_studio_provider_fields_expose_api_base to assert the new contract (`default_value is None`, `placeholder` carries the canonical URL), with a comment pointing at the Greptile threads as the rationale so future contributors don't accidentally re-introduce the default. 26/26 tests in the file pass. JSON validates (`json.load` clean). * feat(azure_ai): add gpt-5.5 to model cost map (#30428) * feat(azure_ai): add gpt-5.5 to model cost map Adds azure_ai/gpt-5.5 and its dated snapshot azure_ai/gpt-5.5-2026-04-23 to both the canonical and bundled cost maps. gpt-5.5 is generally available on Azure AI Foundry; pricing mirrors the openai gpt-5.5 entry, matching the established azure_ai convention (verified identical for gpt-5.4), in the azure tier structure (base / above-272k / priority). supports_minimal_ reasoning_effort is false, the capability that changed from gpt-5.4. Fixes #30306 * Update tests/test_litellm/test_gpt_5_5_model_metadata.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix: guard check_and_fix_namespace against None key (#30435) * fix: guard check_and_fix_namespace against None key When user_id is None, the cache key can be None, causing AttributeError: 'NoneType' object has no attribute 'startswith' in check_and_fix_namespace. Add an early return for None key to prevent the error and the ERROR-level log noise it produces on every unauthenticated request. Fixes #30424 * fix: update type annotations for check_and_fix_namespace - key: str -> Optional[str] (now handles None input) - return: str -> Optional[str] (returns None when input is None) Addresses Greptile review concern about type signature mismatch. * fix: revert check_and_fix_namespace type signature to str to fix MyPy downstream errors * fix: update type annotations for check_and_fix_namespace - Change signature from str -> str to Optional[str] -> Optional[str] - Remove type: ignore comment on None return - Add None guard in async_set_cache_sadd before passing to helper Addresses review feedback from Sameerlite on type mismatch. * Revert "fix: update type annotations for check_and_fix_namespace" This reverts commit5272920fa0. --------- Co-authored-by: michaelxer <michaelxer@users.noreply.github.com> * fix(cost): apply service_tier suffix to above-threshold cache rates and expose priority+threshold keys in ModelInfo (#30450) * fix(cost): apply service_tier suffix to above-threshold cache rates and expose priority+threshold keys in ModelInfo Models that publish both a service_tier (e.g. priority) rate and an above-threshold tier (e.g. _above_200k_tokens) currently bill cached tokens at the standard above-threshold rate rather than the priority above-threshold rate. Affected entries in the live pricing JSON include gemini-3-pro-preview, gemini-3.1-pro-preview and their vertex_ai/ and gemini/ variants, plus azure/gpt-5.4 and azure_ai/gpt-5.4. For a 250K-token priority request with 200K cached tokens against gemini-3-pro-preview, the leak is about 44 percent of the prompt cost. Two stacked defects caused this. First, ModelInfoBase (and the ModelInfo pydantic class) and the get_model_info construction in litellm/utils.py omit the priority+above-threshold cost keys, so even if the calculator asked for them they would never reach it. Second, in _get_token_base_cost the cache_creation/cache_read tiered keys never get wrapped with _get_service_tier_cost_key, while the input/output tiered keys above and below do. The change here surfaces six new keys (input, output and cache_read at both 200k and 272k priority variants) and wraps the three cache tiered keys in _get_token_base_cost the same way input/output already are. _get_cost_per_unit's existing service_tier-to-base fallback covers models that ship the standard above-threshold rate without a priority variant. Adds one regression test in tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py that drives the actual generic_cost_per_token path for gemini-3-pro-preview at 200K cached + 50K text under priority and asserts the priority above_200k rates are picked. Verified the test fails on litellm_internal_staging without these changes and passes with them. * fix(cost): drop guard on cache tiered keys so service_tier fallback can reach standard above-threshold rate Addresses Greptile P1 on PR 30450. The previous commit wrapped cache_creation_tiered_key, cache_creation_1hr_tiered_key, and cache_read_tiered_key with _get_service_tier_cost_key (matching how the sibling input and output tiered keys are wrapped) but kept the surrounding 'if key in model_info' guards. For models that publish a standard above-threshold cache rate but no priority variant (gpt-5.4-pro, gpt-5.5-pro and their dated siblings, plus vertex_ai/claude-sonnet-4-5 for cache_creation), the guard short-circuits before _get_cost_per_unit's existing service_tier-to-base fallback can strip _priority and find the standard above-threshold key. The result on priority requests over the threshold was that those models silently dropped from the above-threshold rate back to the priority-base rate. Dropping the guard and calling _get_cost_per_unit unconditionally (mirroring how tiered_input_key and tiered_output_key are already handled) restores correct billing for that class of models while keeping the new priority+above-threshold behaviour for gemini-3-pro-preview and friends. Adds a second regression test that pins generic_cost_per_token for vertex_ai/claude-sonnet-4-5 priority + above_200k with cached and cache_creation tokens to the expected standard above-threshold rates, so the guard cannot be silently reintroduced for either the cache_read or cache_creation path. * fix(presidio): skip pre-call masking when guardrail is logging_only (#30461) The Presidio pre-call hook masked the live request unconditionally, ignoring the configured event hook. With mode: logging_only the masked request reached the model, so its response echoed anonymization tokens (e.g. <PERSON>) instead of the real output. Gate async_pre_call_hook on should_run_guardrail, matching every other guardrail; logging_only masking still happens via async_logging_hook. * fix(router): resolve list unhashable crash on model alias (#30464) * fix(router): resolve list unhashable crash on model alias Fixes the fallback parsing logic that mistakenly categorized standard array fallback definitions as override dictionaries when a deployment alias matches the literal string 'model'. Closes https://github.com/BerriAI/litellm/issues/30459 * fix(router): address greptile review for fallback parsing edge cases - Resolves ambiguity in standard vs override fallback dictionaries by iterating over all items and validating that no mapped litellm param resolves to a non-list type. - Adds regression tests in test_router_order_fallback.py to prevent unhashable type crash from silently re-entering the codebase. * chore(router): format code with black to pass CI * fix(hosted_vllm): remove thinking_blocks and convert list content to strings (#30475) * fix: hosted_vllm remove thinking_blocks and convert list content to strings vLLM endpoints reject assistant messages with thinking_blocks converted to content list blocks. This change removes thinking_blocks entirely and converts any list content back to strings. This fixes BadRequestError when using Claude Code with hosted_vllm models that pass thinking_blocks in messages. * fix(hosted_vllm): address Greptile review feedback - Join multiple text blocks with newline instead of empty string - Always set content to string (never None) to avoid vLLM validation errors * fix(hosted_vllm): update chat transformation to clean assistant messages * fix: re-raise exception instead of silently dropping MCP team permissions (#30477) * fix: re-raise exception instead of silently dropping MCP team permissions When MCPRequestHandler.get_allowed_mcp_servers raises, the broad except was swallowing the error and returning only allow_all_server_ids, silently discarding all team-level object_permission grants. Fixes #30476 * fix: log full traceback when MCP permission lookup fails Uses verbose_logger.exception() instead of warning() so operators can see the full traceback when team-level object_permission grants are dropped due to an internal error in get_allowed_mcp_servers. Fixes #30476 * fix: remove timezone date expansion in daily-activity aggregation (#29569) * fix: remove timezone date expansion in daily-activity aggregation Single-day spend queries from non-UTC timezones over-counted by ~2x because the previous implementation widened the SQL date range by a full UTC day on whichever side the offset pointed. Spend is bucketed in whole-UTC-day rows in LiteLLM_DailyUserSpend, so the expansion pulled an extra 24h of unrelated bucket data per boundary. Concretely on IST (UTC+5:30, offset -330): a single-day query for 2026-05-29 was rewritten to date >= 2026-05-28 AND date <= 2026-05-29 and returned spend across both UTC days. Sums of single-day queries across a 5-day window then exceeded the equivalent multi-day aggregate by ~50%, which is mathematically impossible. Treat the local date range as the UTC date range. The aggregation table has no hour-level granularity, so any conversion using only date arithmetic must round to whole UTC days; the previous fix turned that boundary slop into systematic over-counting. Pass-through trades a small one-time slop at each end of the range for correct, monotonic, additive results across single-day and multi-day queries. Repro from production: bedrock/global.anthropic.claude-opus-4-8 over 2026-05-29 to 2026-06-02, IST timezone: - 5-day aggregate: $701.39 / 1,831 reqs - Sum of 5 single-day queries: $1,070.94 / 2,755 reqs - Excess (was 1.527x): now matches within boundary slop Adds regression tests in TestAdjustDatesForTimezone and TestBuildAggregatedSqlQuery that pin the pass-through behavior and the additivity invariant for any future implementation. * ci: rerun checks on litellm_oss_branch base --------- Co-authored-by: Sameer Kankute <sameer@berri.ai> * fix: buffer native gemini sse frames (#30225) * fix: buffer native gemini sse frames * fix: scope native gemini sse buffering * fix: check raw sse residual buffer size * feat: updated openrouter provider to map max level to xhigh (#28881) * feat(proxy): allow use_redis_transaction_buffer without redis cache (#28764) * feat(proxy): allow use_redis_transaction_buffer without redis cache * fix(proxy): require host or url for standalone buffer redis * fix(mcp): fail closed when scope filter resolves to no servers (#30353) `_get_allowed_mcp_servers_from_mcp_server_names` returned the caller's full allowed-server set when the requested `mcp_servers` list (path- or header-derived) resolved to nothing. URL/header namespacing therefore appeared to work even when the requested name was unknown or the caller had no grant — `/mcp/<typo>/` silently exposed every server the key could reach. Fail closed instead: when `mcp_servers` is explicitly provided but nothing resolves, return an empty list. The `mcp_servers=None` path (no scope requested) keeps its existing behavior. Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> * fix(token-counter): handle Anthropic tool_reference blocks to stop dropped spend logs (#30302) * fix(token-counter): handle Anthropic tool_reference blocks to stop dropped spend logs `token_counter` did not know about Anthropic tool-search `tool_reference` content blocks, a lightweight pointer to a deferred tool that shows up as `{"type": "tool_reference", "tool_name": ...}`. When such a block appeared in message content, `_count_content_list` fell through to its catch-all branch and raised `Invalid content item type: tool_reference`. On the streaming `anthropic_messages` proxy path that exception nulls `response_cost`, which makes the proxy drop the entire SpendLogs row. The result is a silent cost undercount on any tool-search traffic; the request succeeds for the caller but the spend is never recorded. This adds a `tool_reference` branch that counts the referenced `tool_name` (the full tool definition is already counted via the `tools` param, so only the name is added here) and handles an empty/missing name gracefully. The catch-all error message is updated to list `tool_reference` among the expected types. A regression test asserts that a message containing a `tool_reference` block no longer raises and returns a positive token count, and that an empty `tool_name` is handled without error. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(token-counter): collapse explicit None tool_name to empty string In _count_content_list, c.get("tool_name", "") returns None when the key is present with an explicit None value, and str(None) == "None" which is truthy, causing a spurious token to be counted. Use c.get("tool_name") or "" so both a missing key and an explicit None collapse to an empty string and are skipped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(token-counter): cover catch-all for unknown content block type Adds a regression test that calls `_count_content_list` with an unrecognized content block type and asserts it raises `ValueError` whose message names the offending type and lists `tool_reference` among the supported types. This exercises the previously uncovered catch-all branch (codecov patch gap) and pins the error contract. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(token-counter): cover tool_reference on the spend/cost and streaming paths Adds end-to-end regression tests that exercise the real public entry points (`completion_cost` and `stream_chunk_builder`), not just the private `_count_content_list` helper, for Anthropic tool-search `tool_reference` content blocks. These pin the actual bug the fix addresses: before the fix the `tool_reference` block raised out of `completion_cost` -> the proxy logging layer nulled `response_cost` and the spend callback dropped the SpendLogs row (silent cost undercount on all tool-search traffic); and `stream_chunk_builder` swallowed the same raise and collapsed prompt_tokens to 0. With the fix, cost is positive and prompt_tokens are counted. Verified: 3 fail without the fix, 3 pass with it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(cost): add cost mapping for deepseek-v4-flash and deepseek-v4-pro (#27056) * feat(cost): add cost mapping for deepseek-v4-flash and deepseek-v4-pro Adds pricing entries for the two new DeepSeek V4 models released on 2026-04-24, for both bare model names and the deepseek/ provider prefix. Prices sourced from https://api-docs.deepseek.com/quick_start/pricing: - deepseek-v4-flash: $0.14/M input, $0.28/M output - deepseek-v4-pro: $1.74/M input, $3.48/M output Cache hit price set to 1/10 of input (per DeepSeek docs). Context window: 1M tokens for both models. Closes #26709 * fix(cost): update backup registry for deepseek-v4 * style: remove print statement from deepseek-v4 test * feat(cost): add cost mapping for deepseek-v4-flash and deepseek-v4-pro Adds pricing entries for the two new DeepSeek V4 models released on 2026-04-24, for both bare model names and the deepseek/ provider prefix. Prices sourced from https://api-docs.deepseek.com/quick_start/pricing: - deepseek-v4-flash: $0.14/M input, $0.28/M output - deepseek-v4-pro: $1.74/M input, $3.48/M output Cache hit price set to 1/10 of input (per DeepSeek docs). Context window: 1M tokens for both models. Closes #26709 * fix: update deepseek-v4 prices to active discounted rates * test: update deepseek-v4 prices in tests to match active discounted rates * fix(deepseek): remove duplicate entries and update backup registry to active discounted rates * fix: update max_output_tokens to 384K for deepseek-v4 * fix: correctly restore upstream models accidentally dropped during merge * fix(tests): resolve failing claude-fable-5 and reasoning tests by safely updating cost map - Pulled the latest cost map from upstream staging - Safely appended deepseek-v4 mapping without deleting duplicate keys or formatting via json.dump * fix(tests): correct deepseek model cache prices and update JSON schema - Appended both prefixed and bare deepseek-v4 models to satisfy test assertions - Corrected deepseek-v4-pro expected cache hit and token prices based on latest review updates - Added missing realtime endpoint to test_utils.py INTENDED_SCHEMA * fix: remove accidental azure/gpt-realtime-whisper addition --------- Co-authored-by: Dushyant Acharya <dushyantacharya@Dushyants-MacBook-Pro.local> * feat(key/info): expose per-model budget usage in /key/info response (#30394) * feat(key/info): expose per-model budget usage in /key/info response Add model_max_budget_usage to /key/info and /v2/key/info responses. For each model in model_max_budget, reads current-period spend from the same DualCache used by the budget enforcer and returns it alongside the limit and time period so callers can see how much of each model budget has been consumed in the active window. * test(key/info): add coverage for model_max_budget_usage in v1 and v2 endpoints Add tests for the model_max_budget_usage enrichment in both info_key_fn and info_key_fn_v2, covering the budget-present path, the empty-budget path, and the v2 batch endpoint. * fix(key/info): source model_max_budget current_spend from SpendLogs instead of DualCache The DualCache used for enforcement is ephemeral and only populated when budget metadata is present at request time. Fall back to a direct LiteLLM_SpendLogs DB aggregation using the budget period window (budget_reset_at - budget_duration) for accurate reporting. Also fall back to litellm_budget_table.model_max_budget when the key's top-level field is empty, and round current_spend to 4 decimal places. * test(key/info): cover remaining branches in model_max_budget_usage helpers Add unit tests for: prisma_client=None early return, DB query exception swallowing, invalid budget_duration handled by _compute_budget_period_start, budget_reset_at received as a datetime object (Prisma native type), max_seconds=0 early return, and skipping models that lack a budget_duration. Also remove an unreachable except branch where fromisoformat would fail after _compute_budget_period_start already validated the same value. * test(key/info): cover except path for unparseable per-model budget_duration * fix(key/info): compute per-model rolling windows in model_max_budget_usage Each model in model_max_budget now gets its own time window derived from its own budget_duration, rather than sharing a single window computed as the max (or the budget table's reset_at). This matches what the DualCache enforcer actually tracks and prevents current_spend from being inflated for models with shorter windows. _query_model_spend_for_period is refactored to accept a model filter (handling provider-prefix variants in SQL) and return a float directly. _compute_budget_period_start and the budget_table window path are removed as they are no longer needed. * refactor(model_max_budget_limiter): remove dead get_current_period_spend method * refactor(key/info): strip synthetic formatter noise from PR diff Restore key_management_endpoints.py and test_key_management_endpoints.py to origin/litellm_internal_staging, then re-apply only the intentional additions: _query_model_spend_for_period, _build_model_max_budget_usage, the two endpoint patches (info_key_fn / info_key_fn_v2), and the new test suite. The previous commits had reformatted ~300 pre-existing lines across both files, making the functional diff unreadable. * test(key/info): cover empty-rows path in _query_model_spend_for_period * fix(model_max_budget_limiter): guard BudgetConfig construction inside try/except A malformed model entry in the DB (e.g. non-numeric max_budget from a manually edited or migrated row) caused BudgetConfig(**budget_info) to raise a Pydantic ValidationError outside any exception guard, surfacing as a 500 for the entire /key/info or /v2/key/info call. Merging both try/except blocks into one ensures bad entries are silently skipped, consistent with the existing duration_in_seconds guard. * fix: don't stack provider prefix on wildcard models with a custom prefix (#30360) * fix: don't stack provider prefix on wildcard models with a custom prefix get_known_models_from_wildcard expanded provider-prefixed model ids (e.g. "ollama/gemma3:1b" from get_provider_models) by prepending the wildcard's prefix whenever the id did not already start with it. With a custom wildcard prefix such as "ollama_server1/*" (used to distinguish multiple Ollama instances), this produced "ollama_server1/ollama/gemma3:1b", which is uncallable and breaks /v1/models. When the expanded id already carries a provider prefix, replace it with the wildcard's prefix instead of stacking both. Matching-prefix and bare-model cases are unchanged. Fixes #30358 * fix: only strip a known provider prefix when expanding custom wildcard prefixes The wildcard expansion replaced the leading slash segment of every expanded id with the wildcard prefix whenever the id did not already start with it. For ids whose first segment is an org rather than a litellm provider (for example a provider returning "meta-llama/Llama-3-8B" with no outer provider prefix), that dropped the org and produced an uncallable id Only strip the leading segment when it is a recognized provider (membership in LlmProviders); otherwise keep it and just prepend the wildcard prefix. Provider-prefixed ids like "ollama/gemma3:1b" still have their prefix replaced, so the original fix is unchanged for known providers * address greptile review feedback: log dropped non-text vLLM assistant content blocks (greploop iteration 1) * fix(ci): format credential_form_helpers test + regenerate dashboard schema.d.ts * fix(proxy): raise litellm.BadRequestError for missing model param When no model is passed, route_request now raises a litellm.BadRequestError ('Missing model parameter') instead of falling through to ProxyModelNotFoundError. This keeps the missing-param error clear and independent of router wildcard state. Unknown (non-empty) model names still raise ProxyModelNotFoundError. * Revert "fix(proxy): raise litellm.BadRequestError for missing model param" This reverts commit9240da403c. * Revert "fix(router): clean pattern_router state on upsert/delete (#29601)" This reverts commitad4e6e2395. * fix: correct streaming and key budget usage reporting * fix(hosted_vllm): type assistant tool_calls to satisfy mypy * feat: aws secret manager cross region replication (#30368) * feat(aws-secret-manager): add replica_regions cross-region replication after CreateSecret When store_virtual_keys is enabled, async_write_secret() only wrote secrets to the primary AWS region. Multi-region proxy deployments had no built-in way to synchronize virtual key secrets across regions through LiteLLM, requiring external replication mechanisms. Add replica_regions support to AWSSecretsManagerV2: - New replica_regions field in KeyManagementSettings (types/secret_managers/main.py) - New async_replicate_secret() method that calls ReplicateSecretToRegions API - async_write_secret() calls replication after successful CreateSecret - Replication failure is logged as a warning but does NOT fail key creation - load_aws_secret_manager() forwards replica_regions from key_management_settings Configuration example: key_management_settings: store_virtual_keys: true replica_regions: - us-west-2 - eu-west-1 When replica_regions is omitted or empty, behavior is unchanged. * test(aws-secret-manager): restore litellm.secret_manager_client after test to prevent state pollution * test(aws-secret-manager): add coverage for HTTP error and replication exception paths * fix: restore litellm.secret_manager_client global state in test; add replication log proof - Global state in test_load_aws_secret_manager_passes_replica_regions was already guarded with try/finally (committed in previous pass); no further change needed for Fix 1. - Fix 2: add verbose_logger.info("ReplicateSecretToRegions called …") inside async_replicate_secret so callers get an observable INFO log line whenever replication fires. - Add test_replication_fires_on_create: calls async_replicate_secret directly with caplog.at_level(INFO, logger="LiteLLM") and asserts "ReplicateSecretToRegions" appears in the captured log output, proving the code path executes. * fix: pass request to streaming generators * fix(hosted-vllm): preserve assistant structured content * fix(hosted_vllm): satisfy mypy on preserved structured content assignment * chore: resolve litellm_internal_staging merge conflicts for #30527 (#30554) * chore(codecov): add Batches, Videos, and Realtime components (#30517) * chore(codecov): add Batches, Videos, and Realtime components Define per-feature Codecov components so PR comments track coverage for batch API, video generation, and realtime streaming paths. Co-authored-by: Cursor <cursoragent@cursor.com> * chore(codecov): use wildcard path for Batches proxy component Align batches_endpoints glob with Videos, Realtime, and Proxy_Authentication. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com> * test(batches): move orphan tests into tests/test_litellm for CI coverage (#30510) Four batch-related tests lived under tests/litellm/ and were never picked up by GitHub Actions. Relocate them and fix gemini multimodal e2e to use the batchEmbedContents path expected for gemini/ provider. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(guardrails): run pre_call hook once for model-level guardrails (#30543) * fix(guardrails): run pre_call hook once for model-level guardrails A CustomGuardrail attached to a deployment via litellm_params.guardrails gets its async_pre_call_hook invoked twice per request: once by the proxy pre-call loop and again by async_pre_call_deployment_hook after the router spreads the model-level guardrails into the top-level request kwargs. Record in request metadata that the proxy pre-call loop already ran a given guardrail, and have the deployment hook skip it when the marker is present. Direct-SDK usage never runs the proxy loop, so the deployment hook stays the sole invocation there and still fires exactly once. The marker key is stripped from untrusted caller metadata so a request body cannot suppress a model-only guardrail by pre-seeding it. * fix(guardrails): mark pre_call dedup on the post-hook request data Record the exactly-once marker after async_pre_call_hook runs, on the data object that flows downstream, rather than before it. A guardrail whose hook returns a brand-new request dict (instead of mutating or spreading the one it received) would otherwise discard the marker, letting the deployment hook re-run the guardrail a second time. * fix(guardrails): stop re-initializing DB guardrails on every poll (#30542) * fix(guardrails): stop re-initializing DB guardrails on every poll InMemoryGuardrailHandler._has_guardrail_params_changed compared the in-memory LitellmParams against the raw dict loaded from the DB. The in-memory side carries every field default and coerces enums via model_dump(), while the DB side only holds the keys originally stored, so the two shapes never compared equal and the guardrail was rebuilt on every poll cycle. Each rebuild created a fresh instance, but delete_in_memory_guardrail only removed the old callback from litellm.callbacks. Request handling promotes guardrail callbacks into the success/failure/async lists, so the previous instance stayed referenced there and instances accumulated. Normalize both sides through LitellmParams(...).model_dump() before diffing, and purge the callback from every callback list on delete. * refactor(guardrails): narrow params-normalization fallback to ValidationError The comparison normalizer caught a bare Exception and silently fell back to the raw dict, which hid the cause and quietly degraded the affected guardrail back to re-initializing on every poll. Catch only the ValidationError that LitellmParams construction can raise, log a warning so the offending row is diagnosable, and let any other error surface instead of being swallowed. * refactor(callbacks): add remove_callback_from_all_lists helper to manager Move the knowledge of which callback lists a callback can be promoted into out of the guardrail registry and into LoggingCallbackManager, where the rest of the callback-list bookkeeping already lives. delete_in_memory_guardrail now delegates to the new helper instead of iterating the lists itself. * chore(oss): litellm oss staging 150626 (#30463) * fix(pricing): add GitHub Copilot MAI Code Flash pricing (#30415) * fix(pricing): add GitHub Copilot MAI Code Flash pricing Add GitHub Copilot pricing entries for MAI-Code-1-Flash and the internal Copilot CLI model name so cost calculation can price input, cached input, and output tokens. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(pricing): cover GitHub Copilot MAI Code Flash pricing Add regression coverage for both GitHub Copilot MAI-Code-1-Flash model names, including cached input pricing, chat endpoint metadata, and cost_per_token arithmetic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(router/proxy): propagate completed_response through FallbackResponsesStreamWrapper for streaming /v1/responses container ownership (#30210) (#30213) * fix(router/proxy): propagate completed_response through FallbackResponsesStreamWrapper for streaming /v1/responses container ownership (#30210) #28990 added ownership recording for streaming /v1/responses via _wrap_responses_stream_for_container_ownership, which reads `getattr(stream_response, 'completed_response', None)` to extract the ResponsesAPIResponse. The unit test bypassed the Router, so it never exercised the production wrapping path. Through the Router (every proxy deployment), the stream is wrapped by FallbackResponsesStreamWrapper (router.py:2527). Its __init__ set `self.completed_response = None` and __anext__ only forwarded chunks — the inner source iterator's terminal event never bubbled up to the attribute the ownership hook reads, so the hook silently recorded nothing and every follow-up /v1/containers/<id>/files call returned 403 for non-admin keys. This commit: - router.py: pre-resolves the responses-API terminal event tuple (response.completed / .incomplete / .failed) once per _aresponses_streaming_iterator call, and has the wrapper's __anext__ sniff each forwarded chunk's .type. First terminal event hit gets stored on the wrapper's completed_response. Iterator-agnostic — works for source_iterator AND any future wrapper. - common_request_processing.py: when _extract_completed_responses_response returns None we now warn instead of silently skipping. Reporter on #30210 lost a day to this exact silent skip; the warning surfaces future regressions of the same shape directly in operator logs. Fixes #30210 * fix(router): type-ignore wrapper getattr-defaults; broaden ownership-skip warning CI lint (mypy) flagged the three pre-existing getattr(..., None) assignments in FallbackResponsesStreamWrapper.__init__: router.py:2564 self.response = getattr(source_iterator, 'response', None) router.py:2565 self.model = getattr(source_iterator, 'model', None) router.py:2566 self.logging_obj = getattr(..., None) Those lines also exist on litellm_internal_staging and pass mypy there. Adding the typed terminal-event tuple above the class made the function body more narrowable, which surfaced the pre-existing mismatch — base class declares non-Optional types but the bridge path (LiteLLMCompletionStreamingIterator) legitimately omits these. Keep the None fallback and silence with type: ignore[assignment]. Greptile 4/5 note: the ownership-skip warning hard-named code_interpreter which misleads operators when a non-code_interpreter stream aborts. Generalize to 'any tool container (e.g. code_interpreter)'. * fix(register_model): drop synthesized zero costs to preserve sparse entries (#30198) (#30201) * fix(register_model): drop synthesized zero costs to preserve sparse entries (#30198) get_model_info synthesizes input_cost_per_token / output_cost_per_token = 0 when they are absent from the raw entry (the price-unknown and free cases share the same representation). register_model then merges that result back into litellm.model_cost, which flips a sparse entry from 'no cost keys' (priced via model name) to 'cost keys = 0' (free). That defeats _is_cost_explicitly_configured (#24949) on re-registration: _is_model_cost_zero returns True, common_checks skips every tag / key / team / user / org budget check for the group, and over-budget traffic keeps returning 200. Spend keeps recording because cost calc still resolves by model name, so the symptom is silent and only triggers on the second register_model pass (router rebuild, /model/update, config sync). Mirror the existing litellm_provider-None guard one block above and pop the cost fields from the synthesized result when they are absent from the raw entry and not in the caller's value. Caller-provided zeros (genuinely free models, BYOK overrides) are preserved. Fixes #30198 * fix(register_model): switch _raw_entry to is-None checks + drop dead test assertion Greptile #30201 review notes: - the `or`-chain in the raw-entry lookup treated an empty dict (a key with no fields) as falsy and fell through to the second arm — replace with explicit `is None` checks so a present-but-empty entry is still taken at face value. - the first assertion in `test_router_double_init_keeps_db_model_entry_sparse` used `in (None, 0)` which passes under the bug condition (cost = 0 matches the tuple); the strong follow-up assertion already covers every shape, so drop the dead branch. * fix(bedrock mantle): use unique function-call id for responses->chat tool calls (#30426) * fix(bedrock mantle): use unique function-call id for responses->chat tool calls ... * fix(bedrock mantle): scope unique tool-call id fallback to degenerate call_id The previous revision preferred the Responses item id for every tool call, which broke providers (and existing tests) where call_id is a unique, canonical correlation key. Restrict the fallback to the degenerate index-based call_id that Bedrock Mantle returns (call_0, call_1, ... resetting per response) and keep call_id otherwise. Revert the change to the OUTPUT_ITEM_DONE streaming handler, whose tool_call_chunk is never emitted (dead code, per review). Extend the regression tests to assert a normal call_id is preserved. * fix(router): preserve azure_ad_token through CredentialLiteLLMParams for /v1/files + batches (#30235) (#30241) * fix(router): preserve azure_ad_token through CredentialLiteLLMParams for /v1/files + batches (#30235) Router.get_deployment_credentials_with_provider re-validates a deployment's litellm_params through CredentialLiteLLMParams before handing them to file/batch/passthrough callers: return CredentialLiteLLMParams( **deployment.litellm_params.model_dump(exclude_none=True) ).model_dump(exclude_none=True) Any field NOT declared on CredentialLiteLLMParams gets silently dropped on the way through. azure_ad_token was undeclared, so Azure deployments using OAuth/M2M (azure_ad_token instead of a static api_key) silently lost their token at the files endpoint and the proxy returned: Missing credentials. Please pass one of api_key, azure_ad_token, azure_ad_token_provider, ... Declare azure_ad_token on CredentialLiteLLMParams alongside api_key / api_base / api_version so it rides through the round-trip. Static-key deployments stay unaffected (Optional, default None, dropped by exclude_none=True). Provider-callable (azure_ad_token_provider) is a separate concern and out of scope here. Fixes #30235 * fix(ui-types): regenerate schema.d.ts for new azure_ad_token field CI's 'Verify schema.d.ts matches the proxy OpenAPI spec' check auto-detected the new field and emitted the exact diff to apply. Two schemas had `aws_secret_access_key` from CredentialLiteLLMParams, both get the new azure_ad_token marker next to it. * fix(proxy): org_admin with own user_id now sees all org teams on /v2/team/list (#30247) When the UI sends the callers own user_id (as it does for non-Admin global roles), _enforce_list_team_v2_access now nulls it out for org admins so _build_team_list_where_conditions scopes by organization_id only -- matching the legacy /team/list behavior and the documented intent. Fixes #30215 Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * test(vertex_ai): multi-region regression coverage for cachedContents host (#29571) (#29707) litellm_internal_staging already routes the cachedContents URL through get_vertex_base_url, fixing the multi-region 404 reported in #29571 — but carries no test coverage for the actual regression scenario (eu/us must resolve to the REP host aiplatform.{geo}.rep.googleapis.com). Add TestContextCachingMultiRegionUrls: parametrized eu/us REP-host assertions (including absence of the old broken {geo}-aiplatform host), plus regional (us-central1) and global no-regression checks. * fix(proxy): close upstream LLM stream when client disconnects mid-stream (#30245) * fix(proxy): close upstream LLM stream when client disconnects mid-stream When a streaming client disconnects, Starlette abandons the response body iterator without calling aclose(), so the proxy's connection to the upstream backend stays open until garbage collection, which may never come. The backend (e.g. vLLM) keeps generating into a dead pipe: small responses drain invisibly into TCP buffers while large ones block the backend on a full send buffer indefinitely (observed via lsof as an ESTABLISHED proxy->backend connection minutes after the client left) create_response now returns a StreamingResponse subclass that closes both its body iterator and the wrapped upstream-facing generator in a shielded finally. The upstream generator is closed directly rather than through a cascade because aclose() on a never-started generator skips its body, which would make the cascade a no-op when the client disconnects before the first chunk is sent. async_streaming_data_generator also gains the same shielded finally-aclose that async_data_generator in proxy_server.py already had, covering the Anthropic and Google SSE paths With this, killing a streaming client causes the backend to observe the abort within about a second and free its slot, while completed streams are unaffected. No flag is needed, unlike the non-streaming opt-in cancel in #30223: this only releases resources after the client is already gone and does not change any response a client can observe Fixes #30244 * fix(proxy): close upstream even when body iterator aclose raises BaseException Addresses the Greptile finding on #30245: the cleanup loop caught only Exception while the generator-level cleanup catches BaseException, so a CancelledError or GeneratorExit escaping body_iterator.aclose() would skip closing the upstream generator. Both sites now use the same scope and a regression test pins that the upstream is closed even when the body iterator explodes with a BaseException * fix(llms): expose aclose on BaseModelResponseIterator so stream close reaches the provider connection The response-level close added for #30244 only worked for SDK-based providers (e.g. openai), whose streams expose aclose all the way down. Providers served by base_llm_http_handler (hosted_vllm and most modern transformation-based providers) wrap a bare response.aiter_lines() generator in BaseModelResponseIterator, which had no aclose or close at all, and nothing retained the httpx response object; so CustomStreamWrapper.aclose() silently did nothing and the upstream connection stayed open. Verified with a vLLM-style mock: with hosted_vllm/ the backend streamed all 100 chunks to completion after the client disconnected, while openai/ aborted at chunk 6 BaseModelResponseIterator now carries an optional http_response and an aclose() that closes it; make_async_call_stream_helper attaches the response after building the iterator. With this, hosted_vllm aborts the backend within ~1.6s of the client dropping, and completed streams are unaffected --------- Co-authored-by: kursad <kursad.lacin@brado.net> * feat(anthropic): surface compaction usage iterations data (#27065) * feat(anthropic): surface compaction usage iterations data * style: apply black formatting to fix lint checks * fix(usage): correct calculate usage with cached tokens when use ChatCompletionUsageBlock (#30422) * fix(usage): correct calculate usage with cached tokens when use ChatCompletionUsageBlock * fix(usage): optimize test imports * feat: add fastCRW search provider (#30434) * feat(provider): add LibertAI as a JSON-configured OpenAI-compatible provider (#30203) * feat(provider): add LibertAI as a JSON-configured OpenAI-compatible provider * libertai: update served endpoints backup + add mode/matrix tests Addresses review feedback: - Add libertai to litellm/provider_endpoints_support_backup.json, the file actually served by GET /public/supported_endpoints (the root provider_endpoints_support.json already had it). - Add tests asserting bge-m3 normalizes to mode='embedding' and that the served matrix lists libertai. embeddings stays false: the JSON-configured provider path only wires chat routing (OpenAILike embedding handler is reached only for literal openai_like/llamafile/lm_studio), matching the llamagate precedent; bge-m3 remains in the cost map for metadata. --------- Co-authored-by: Moshe Malawach <moshemalawach@users.noreply.github.com> * feat(provider): add ModelScope as an OpenAI-compatible provider (#28460) * add ModelScope API support * add modelscope api support * update modelscope model list * add image-genetation support * update test and multimodal * fix: address PR review feedback for modelscope provider * update README * fix(customer_endpoints): restrict /customer/daily/activity to admin-only (#28849) * fix(customer_endpoints): restrict /customer/daily/activity to admin-only * fix(customer_endpoints): check role before prisma_client guard * fix(custom_guardrail): key disable_global_guardrails takes precedence over team guardrail list (#28563) * fix(fallbacks): preserve fallback model in SDK fallback responses (#28260) * fix(fallbacks): preserve fallback model in response when using SDK-level fallbacks * fix(fallbacks): gate x-litellm-* passthrough to trusted callers only The previous patch unconditionally let `x-litellm-*` keys bypass the `llm_provider-` prefix in `process_response_headers`. That function is also called on raw upstream-provider response headers (e.g. from `llm_http_handler.py`), so a malicious provider could return `x-litellm-attempted-fallbacks` and spoof a LiteLLM-internal marker, bypassing the proxy model-override guard. Add a `preserve_litellm_internal_headers` flag (default False). Only `response_metadata.py`, which re-processes the already-built `_hidden_params["additional_headers"]` dict (LiteLLM-owned), passes True. Raw provider header callsites keep the default False, so upstream `x-litellm-*` still gets the `llm_provider-` prefix. Adds a regression test for the spoofing case and renames the existing preserve test to make the trusted-path semantics explicit. * fix(fallbacks): ignore preserve_litellm_internal_headers for raw httpx.Headers inputs * style(core_helpers): apply black formatting * fix(lint): remove banned typing.List/Dict/Any imports and suppress PLR0913 on interface overrides Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): apply black formatting to modelscope chat transformation Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): replace noqa with proper fixes — use **kwargs and Awaitable instead of Any/List Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): remove unused AllMessageValues import Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * revert: restore base_model_iterator.py to original PR state Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): restore full method signatures for MyPy compatibility; bump PLR0913 budget for new provider files Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(lint): use @override to suppress PLR0913 on inherited signatures instead of bumping budget The overrides keep their full base-class signatures for MyPy compatibility, but those signatures carry more than five parameters, which tripped PLR0913 on each subclass redeclaration. Since the arity is dictated by the base class and cannot be reduced, decorate the overrides with typing_extensions.override; ruff treats that as the intended signal that the parameter count is not under the author's control and skips PLR0913. This restores the PLR0913 baseline to 1813. * fix(lint): add @override to modelscope image generation overrides Apply the same typing_extensions.override treatment to the image generation config so its inherited-signature overrides do not count against PLR0913. --------- Co-authored-by: Joel Tony <github@jaytau.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: hcl <chenglunhu@gmail.com> Co-authored-by: ztko <96878659+koztkozt@users.noreply.github.com> Co-authored-by: Nahrin <nahrin@nahrinoda.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Humphrey <a739376838@gmail.com> Co-authored-by: kursadlacin <kursadlacin@gmail.com> Co-authored-by: kursad <kursad.lacin@brado.net> Co-authored-by: Dushyant Acharya <dushyantacharya873@gmail.com> Co-authored-by: Yuriy <yuriy.shuyskiy@gmail.com> Co-authored-by: Recep S <22618852+us@users.noreply.github.com> Co-authored-by: Moshe Malawach <moshe.malawach@protonmail.com> Co-authored-by: Moshe Malawach <moshemalawach@users.noreply.github.com> Co-authored-by: Rongkun Yan <2493404415@qq.com> Co-authored-by: Varshith <kvarshithgowda@gmail.com> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> * ci(lint): add blanket-noqa, dataclass-default, and unused-noqa Ruff rules (#30516) * ci(lint): enforce blanket-noqa, dataclass-default, and unused-noqa rules Enable PGH004 (blanket-noqa), RUF008 (mutable-dataclass-default), RUF009 (function-call-in-dataclass-default-argument), and RUF100 (unused-noqa) in ruff.toml, and clean up every resulting violation. RUF008/RUF009 were already clean. PGH004/RUF100 surfaced ~335 stale or blanket noqas: blanket `# noqa` are now scoped to the rule they actually suppress (mostly T201), dead directives are removed, and inapplicable codes are trimmed (e.g. F401 dropped from `import *`). lint.external lists rules enforced outside this config (the strict-rule gate via ruff-strict.toml and upstream litellm's own ruff config) so RUF100 keeps the noqa directives that protect them instead of stripping coverage this config can't see. * ci(lint): trim RUF100 external list to load-bearing codes only Drop the 9 precautionary strict-gate codes (ANN001/002/003/401, B006, PLR0913, PLW0603, RUF012, TID251) that have zero `# noqa` references in the gated source. Keep only the 11 codes with live suppressions so RUF100 doesn't flag them as unused. Future strict-gate suppressions can re-add codes here (or fix the underlying issue) as needed. * ci: ratchet lint and type-check gates (ruff preview, ANN, mypy, basedpyright) (#30379) * ci: enable ruff preview rules under the budgeted strict gate Turn on ruff preview in the strict-budget lane (ruff-strict.toml) only, leaving the clean gate (ruff.toml) untouched so make lint-ruff stays at zero. Enumerate the 118 firing codes explicitly with explicit-preview-rules so the gate is deterministic and stable across ruff upgrades rather than depending on preview auto-selecting the broad catalog. Grandfather the existing 58438 violations into ruff-strict-budget.json as per-rule baselines with headroom, so only net-new violations fail CI. The existing ten rules keep their hand-tuned slack; the new rules get slack 10 when the baseline is 50 or more and 3 otherwise. * ci: add ANN return-type rules to the budgeted strict gate Add ANN201/202/204/205/206 (missing return annotations) to the strict lane and grandfather the existing counts into ruff-strict-budget.json so the codebase ratchets toward explicit return types without breaking CI. * ci: add mypy (disallow_untyped_defs) and basedpyright strict gates with baselines Add two type-check gates, each grandfathering the current tree so only net-new violations fail CI, matching the ruff strict-budget ratchet. mypy gains disallow_untyped_defs in litellm/mypy.ini (the config the CI invocation actually reads; the root [tool.mypy] is not picked up from the litellm/ working dir). The 4885 existing missing-annotation errors are captured in litellm/.mypy-baseline.txt and the run is piped through mypy-baseline filter so new untyped defs are rejected. basedpyright runs in strict mode over litellm/, with enableTypeIgnoreComments disabled so it only honors '# pyright: ignore' and never polices mypy's '# type: ignore'. The existing strict diagnostics are grandfathered into .basedpyright/baseline.json. Both tools are pinned in the dev group and uv.lock; the lint workflow and Makefile run them filtered through their baselines, with lint-mypy-baseline-update and lint-basedpyright-baseline-update to ratchet. * ci: raise lint job timeout to 15m for the basedpyright strict pass * ci: pin pythonVersion 3.12 and regenerate baselines against merged base Merge litellm_internal_staging so the baselines cover code the CI merge includes (e.g. the cisco_ai_defense guardrail), which otherwise tripped the mypy gate with 3 ungrandfathered no-untyped-def errors. Pin pythonVersion 3.12 in pyrightconfig so basedpyright's strict analysis is reproducible across interpreter versions (CI runs 3.12). * ci: regenerate basedpyright baseline against the frozen lint env The previous baseline was generated with optional provider deps (azure, google, anthropic, mcp, numpydoc, google-genai) installed locally, so CI's dev-only env surfaced ~3500 reportUnknown*/reportMissingTypeStubs errors not in the baseline. Regenerate after uv sync --frozen so the baseline reflects the same dependency set the lint job sees. * ci: regenerate basedpyright baseline on python 3.12 frozen env The prior baseline still carried proxy-dev packages (e.g. prisma) that the lint job's dev-only, python 3.12 env lacks, leaving 2 unresolved-import errors ungrandfathered. Regenerate in a python 3.12 venv synced to the frozen lock with default groups only, so the baseline matches exactly what CI sees. * ci: replace type-check baselines with per-file count budgets The mypy and basedpyright baselines were position-sensitive (and the basedpyright one was a 27MB file), so ordinary line shifts churned them. Replace both with a per-file count gate: scripts/type_check_gate.py reduces each tool's output to errors-per-file and checks it against a committed {file: max} budget, ignoring line and column numbers. A file fails only when it gains more errors than its ceiling; debt can't be shuffled between files because each file has its own cap and new files default to zero. Budgets (mypy-file-budget.json 48K, basedpyright-file-budget.json 96K) are generated in the python 3.12 frozen lint env so they match CI. Drops the mypy-baseline dependency; basedpyright runs without its native baseline. ratchet via make lint-mypy-budget-update / lint-basedpyright-budget-update. * ci: add a small per-file slack to the type-check gate Allow each file to drift PER_FILE_SLACK (5) errors past its recorded count before failing, so a basedpyright inference ripple in an unrelated file doesn't break the build over a couple of errors. Budgets still record exact counts; the tolerance is applied at check time. * ci: move type-check slack into the budget json and trim lint timeout Make slack declarative: the budget is now {"slack": N, "files": {path: count}} so the tolerance is tuned in JSON without editing the script, mirroring how ruff-strict-budget.json carries its slack. --update preserves the existing slack. Also drop the lint job timeout from 15m to 10m; the mypy and basedpyright passes add ~2m, leaving the job around 4-5m, so 10m is a comfortable margin. * ci: collapse fully-adopted ruff categories and drop inert preview flag ANN (all nine non-removed rules) and BLE (its only rule) were spelled out code-by-code; replace each with its category selector, which is exactly equivalent in 0.15.3 (the removed ANN101/ANN102 are skipped by a category selector and error when named explicitly). explicit-preview-rules was inert: every selected rule is stable and nothing is selected by category, so the flag had nothing to gate. Verified the strict-rule counts are identical before and after (62379 each, zero per-rule drift), so no budget change. * ci: drop redundant pyright dev dependency Nothing invokes bare pyright in the Makefile, the linting workflow, or scripts; the basedpyright gate added on this branch is the only type checker that runs. basedpyright is a superset fork that reads the same pyrightconfig.json and honors the same "# pyright: ignore" comments, so pyright==1.1.408 in the ci group was dead weight. Regenerated uv.lock under the same exclude-newer cutoff so the only change is removing pyright and its package stanza * ci: un-weaken mypy and error on Any in basedpyright mypy: enable warn_return_any, drop the valid-type silencer, and stop globally ignoring missing first-party imports via [mypy-litellm.*] ignore_missing_imports = False, which surfaced eight real broken litellm.* imports the blanket ignore was hiding; third-party imports stay ignored. The per-file budget moves 4888 -> 5799 (902 no-any-return, 1 valid-type, 8 import-not-found), all grandfathered so only net-new errors fail and the ceilings ratchet down basedpyright: error on reportExplicitAny and reportAny. The per-file budget moves 117033 -> 148946 (6931 explicit-Any, 24954 Any-typed expressions), grandfathered the same way * ci: add Any-discipline gate on changed lines under litellm/ Add scripts/check_any_discipline.py, a type-aware gate that fails when a changed line holds a value typed Any -- including the X | Any unions that mypy --strict / basedpyright accept (e.g. re.Match.group() -> str | Any, json.loads() -> Any, bare dict -> dict[Any, Any]). It reuses the repo's mypyc-compiled mypy 1.19 via a custom generic AST walker (mypyc precludes subclassing TraverserVisitor), loads litellm/mypy.ini for parity with lint-mypy, and uses a dedicated incremental cache (.mypy_cache_any) with mtime+hash invalidation to force re-checks. Scope is changed-lines-only so editing a legacy file never forces cleaning its existing Any debt; suppress a genuine typed/untyped boundary with # any-ok: <reason> (ANY002 requires the reason). Wire it into the Makefile (lint-any, lint, lint-dev), a parallel any-discipline CI job with its own actions/cache, .gitignore, and the CLAUDE.md / CONTRIBUTING.md docs. * ci: move Any-gate codes into the shared LIT namespace Renumber the Any-discipline checker into the LIT*** scheme owned by scripts/check_type_discipline.py (PR #30500) so the two checkers share one rule namespace and suppression convention: ANY001 -> LIT002 (Any-typed value; LIT002 was the retired/free slot) ANY002 -> LIT005 (any-ok without a reason; the shared suppression-reason code) ANY000 -> LIT000 (setup/build/read error; the shared error code) Messages and behavior are unchanged; LIT005's text already matches the "<token> requires a reason" shape used for cast-ok/guard-ok. * ci: gate mypy and basedpyright per error rule, not per file Switch the mypy/basedpyright budget gate from per-file error counts to per-rule-code totals, mirroring the {rule: {baseline, slack}} shape of ruff-strict-budget.json. A rule fails when its codebase-wide error count exceeds baseline + slack, so violations are tracked by category rather than by file location. scripts/type_check_gate.py now parses mypy from its text output (trailing [code]) and basedpyright from --outputjson (the JSON `rule` field), since basedpyright's wrapped text diagnostics mis-attribute the rule on continuation lines. Replace the *-file-budget.json files with freshly captured *-code-budget.json baselines and update the Makefile, CI, and CLAUDE.md accordingly. * docs: prefer Pydantic validation over any-ok suppression Point the Any-discipline guidance at validating Any with Pydantic (a model or TypeAdapter that returns a typed value or raises) and frame # any-ok as a last resort that should ideally never be used. * chore: remove extraneous comment * chore: make the CLAUDE.md more concise * chore: clean up bloated CONTRIBUTING.md additions * chore: make Makefile more concise * ci: add the lint-budget-update target CLAUDE.md references CLAUDE.md tells contributors to run make lint-budget-update, but the target was never defined. Add it as an aggregate that re-captures the ruff, mypy, and basedpyright budgets in one shot. * ci: recapture mypy and basedpyright budgets in the lint env The per-rule baselines were captured in a richer dependency env than the CI lint job's uv sync --frozen, so CI resolved fewer types and reported more errors than the budgets allowed (no-any-return 902 over cap 900, plus several basedpyright reportUnknown* rules). Regenerate both in the frozen env so they grandfather the true CI debt: mypy 5786 -> 5799 (no-any-return 890 -> 902, valid-type 1 restored), basedpyright 146213 -> 148942. * ci: check out PR head sha in lint and any-discipline jobs The default pull_request checkout uses refs/pull/N/merge, which folds the latest base commits into HEAD. The diff-based gates (ruff delta, Any discipline) then diff against the event's older base.sha and blame base's own new commits on this branch; staging's otel-v2 and streaming changes (#30326, #30485) tripped the Any gate on files this branch never touched. Checking out the PR head sha makes the gates diff the real branch tip against base, and pins the tree the mypy/basedpyright budgets were captured against so their counts stay deterministic as the base advances. * ci(lint): renumber Any-typed-value rule LIT002 -> LIT009 Free up LIT002 for the sibling type-discipline gate (check_type_discipline.py, #30500), which groups its mutable-collection family at LIT001 (annotation) and LIT002 (construction). This gate's Any-typed-value rule moves to LIT009 so the shared LIT namespace stays contiguous with no holes; LIT000 and LIT005 are unchanged. * style: rename lint-strict-budget -> lint-ruff-budget * ci: harden type-check gates against silent passes (greptile review) type_check_gate.py: refuse to certify a vacuous run. The CI pipe swallows the tool's exit code ('tool || true'), so a crashed mypy/basedpyright that emits nothing would parse to zero errors, breach no ceiling, and pass. is_vacuous_run() now fails when nothing was parsed but the budget expects errors. Also wrap basedpyright's json.loads in a JSONDecodeError handler that prints the offending output instead of dumping a raw traceback. check_any_discipline.py: ALL_LINES was None, which dict.get() also returns for a path absent from the line map, so a path-normalisation mismatch could let a violation on an unchanged file pass the scope filter. Make ALL_LINES a distinct sentinel object so 'whole file' and 'path missing' are unambiguous. Adds tests for all three. --------- Co-authored-by: Sameer Kankute <sameer@berri.ai> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Yassin Kortam <yassin@berri.ai> Co-authored-by: Joel Tony <github@jaytau.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: hcl <chenglunhu@gmail.com> Co-authored-by: ztko <96878659+koztkozt@users.noreply.github.com> Co-authored-by: Nahrin <nahrin@nahrinoda.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Humphrey <a739376838@gmail.com> Co-authored-by: kursadlacin <kursadlacin@gmail.com> Co-authored-by: kursad <kursad.lacin@brado.net> Co-authored-by: Dushyant Acharya <dushyantacharya873@gmail.com> Co-authored-by: Yuriy <yuriy.shuyskiy@gmail.com> Co-authored-by: Recep S <22618852+us@users.noreply.github.com> Co-authored-by: Moshe Malawach <moshe.malawach@protonmail.com> Co-authored-by: Moshe Malawach <moshemalawach@users.noreply.github.com> Co-authored-by: Rongkun Yan <2493404415@qq.com> Co-authored-by: Varshith <kvarshithgowda@gmail.com> * chore: satisfy strict-rule and any-discipline gates for the staging bundle The strict-rule budget and any-discipline gates added in #30379 flag the bundle's new lines: blind-except (BLE001), legacy typing imports (UP006/UP035/UP045), and values typed Any on changed lines (LIT009). Type-fix the cleanly-fixable cases (function signatures, payload dicts as dict[str, object], BudgetConfig.model_validate over **kwargs, direct KeyManagementSettings attribute access over getattr, Optional[X] -> X | None) and suppress the irreducible untyped boundaries (request/streaming dicts, cache reads, httpx responses, asyncio primitives, Pydantic model_dump navigation) with # any-ok and a short reason. Also fix two any-discipline gate false positives so legitimate code is no longer flagged: the synthetic Any in Coroutine/Generator send and yield protocol slots (the awaited/returned value is still checked), and the special-form Any of a TypedDict field's TempNode rvalue placeholder. * chore: extend basedpyright slack to the two rules #30563 left at default PR #30563 raised basedpyright slack to ~10% of baseline across the noisy reportUnknown*/reportAny family so staging bundles clear the per-rule gate, but it left reportArgumentType (slack 3) and reportPrivateUsage (slack 10) at their original tight values. This bundle pushes those two 10 and 1 over their caps respectively, so apply the same ~10% policy: reportArgumentType baseline 1863 -> slack 180, reportPrivateUsage baseline 1625 -> slack 160. No baselines move; only the slack on these two rules * fix: handle duplicate tool calls and stream tail disconnects * fix(proxy): mark stream completed before tail yields, not after [DONE] Clients routinely close the connection right after the final chunk or the terminating data: [DONE] frame. Setting stream_completed only after those trailing yields made the GeneratorExit from that close fall into the disconnect branch, recording false 499 client_disconnected metadata for a response that already delivered all content and fired success logging, and double-releasing the max_parallel_requests slot the success callback had already released. Restore stream_completed before the trailing raw-SSE, error, and [DONE] yields so terminal-marker closes are treated as the successful completions they are. The tool_use dedupe guard is kept. --------- Co-authored-by: apshada <49001649+apshada@users.noreply.github.com> Co-authored-by: Aarkin Karnik <56022539+Aarkin7@users.noreply.github.com> Co-authored-by: David Bochenski <david@goincremental.com> Co-authored-by: Cai Songrui <1922909737@qq.com> Co-authored-by: Martin Honermeyer <7229+djmaze@users.noreply.github.com> Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: fangkang <fangkangm@gmail.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Michael <52305679+michaelxer@users.noreply.github.com> Co-authored-by: michaelxer <michaelxer@users.noreply.github.com> Co-authored-by: Anuj ojha <ojhaanuj224@gmail.com> Co-authored-by: 安妮的心动录 <74543653+anneheartrecord@users.noreply.github.com> Co-authored-by: Zekeriya Akgül <zkry.akgul@gmail.com> Co-authored-by: Thomas Menard <menardorama@gmail.com> Co-authored-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> Co-authored-by: Emerson Gomes <emerson.gomes@thalesgroup.com> Co-authored-by: Mark Lopez <m@silvenga.com> Co-authored-by: Varshith <kvarshithgowda@gmail.com> Co-authored-by: Huynh Duc Tran <110240973+hdt12a1@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: Samarth Maganahalli <samarth.maganahalli@rubrik.com> Co-authored-by: Dushyant Acharya <dushyantacharya873@gmail.com> Co-authored-by: Dushyant Acharya <dushyantacharya@Dushyants-MacBook-Pro.local> Co-authored-by: Thijmen Stavenuiter <thijmenstavenuiter@gmail.com> Co-authored-by: Vineeth Sai <vineethsai4444@gmail.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: rvishwas26 <rvishwas@athenahealth.com> Co-authored-by: Yassin Kortam <yassin@berri.ai> Co-authored-by: Joel Tony <github@jaytau.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: hcl <chenglunhu@gmail.com> Co-authored-by: ztko <96878659+koztkozt@users.noreply.github.com> Co-authored-by: Nahrin <nahrin@nahrinoda.com> Co-authored-by: Humphrey <a739376838@gmail.com> Co-authored-by: kursadlacin <kursadlacin@gmail.com> Co-authored-by: kursad <kursad.lacin@brado.net> Co-authored-by: Yuriy <yuriy.shuyskiy@gmail.com> Co-authored-by: Recep S <22618852+us@users.noreply.github.com> Co-authored-by: Moshe Malawach <moshe.malawach@protonmail.com> Co-authored-by: Moshe Malawach <moshemalawach@users.noreply.github.com> Co-authored-by: Rongkun Yan <2493404415@qq.com>
This commit is contained in:
Sameer Kankute
GitHub
parent
cd26f7d77a
commit
1ccc1e5b23
66 changed files with 4875 additions and 229 deletions
|
|
@ -5,7 +5,7 @@
|
|||
},
|
||||
"reportArgumentType": {
|
||||
"baseline": 1863,
|
||||
"slack": 3
|
||||
"slack": 180
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"baseline": 220,
|
||||
|
|
@ -113,7 +113,7 @@
|
|||
},
|
||||
"reportPrivateUsage": {
|
||||
"baseline": 1625,
|
||||
"slack": 10
|
||||
"slack": 160
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"baseline": 8,
|
||||
|
|
|
|||
|
|
@ -369,6 +369,8 @@ class RedisCache(BaseCache):
|
|||
"""
|
||||
Make sure each key starts with the given namespace
|
||||
"""
|
||||
if key is None:
|
||||
return key # type: ignore[return-value]
|
||||
if self.namespace is not None and not key.startswith(self.namespace):
|
||||
key = self.namespace + ":" + key
|
||||
|
||||
|
|
|
|||
|
|
@ -510,6 +510,8 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv(
|
|||
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
|
||||
)
|
||||
|
||||
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED = 499
|
||||
|
||||
EMAIL_BUDGET_ALERT_TTL = int(
|
||||
os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)
|
||||
) # 24 hours in seconds
|
||||
|
|
|
|||
|
|
@ -5446,6 +5446,39 @@ class StandardLoggingPayloadSetup:
|
|||
error_rate_limit_type=rate_limit_type,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_error_information_for_logging_payload(
|
||||
metadata: dict,
|
||||
original_exception: Exception | None,
|
||||
error_str: str | None,
|
||||
) -> tuple[StandardLoggingPayloadErrorInformation, str | None]:
|
||||
error_information = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=original_exception,
|
||||
)
|
||||
if not metadata.get("client_disconnected"): # any-ok: untyped metadata
|
||||
return error_information, error_str
|
||||
|
||||
client_disconnect_error = metadata.get( # any-ok: untyped metadata
|
||||
"error_information"
|
||||
)
|
||||
if isinstance(client_disconnect_error, dict): # any-ok: untyped metadata
|
||||
error_information = cast(
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
client_disconnect_error, # any-ok: untyped metadata
|
||||
)
|
||||
else:
|
||||
error_information = cast(
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
{ # any-ok: untyped metadata
|
||||
"error_code": "499",
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
},
|
||||
)
|
||||
if not error_str:
|
||||
error_str = "Client disconnected the request"
|
||||
return error_information, error_str
|
||||
|
||||
@staticmethod
|
||||
def get_response_time(
|
||||
start_time_float: float,
|
||||
|
|
@ -5773,8 +5806,12 @@ def get_standard_logging_object_payload(
|
|||
api_base=litellm_params.get("api_base"),
|
||||
)
|
||||
|
||||
error_information = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=original_exception,
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata=metadata, # any-ok: untyped metadata
|
||||
original_exception=original_exception,
|
||||
error_str=error_str,
|
||||
)
|
||||
)
|
||||
|
||||
## get final response object ##
|
||||
|
|
|
|||
|
|
@ -303,40 +303,54 @@ def _get_token_base_cost(
|
|||
|
||||
# Apply tiered pricing to cache costs
|
||||
cache_creation_tiered_key = (
|
||||
f"cache_creation_input_token_cost_above_{threshold_str}_tokens"
|
||||
_get_service_tier_cost_key(
|
||||
f"cache_creation_input_token_cost_above_{threshold_str}_tokens",
|
||||
service_tier,
|
||||
)
|
||||
if service_tier
|
||||
else f"cache_creation_input_token_cost_above_{threshold_str}_tokens"
|
||||
)
|
||||
cache_creation_1hr_tiered_key = (
|
||||
_get_service_tier_cost_key(
|
||||
f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens",
|
||||
service_tier,
|
||||
)
|
||||
if service_tier
|
||||
else f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens"
|
||||
)
|
||||
cache_creation_1hr_tiered_key = f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens"
|
||||
cache_read_tiered_key = (
|
||||
f"cache_read_input_token_cost_above_{threshold_str}_tokens"
|
||||
_get_service_tier_cost_key(
|
||||
f"cache_read_input_token_cost_above_{threshold_str}_tokens",
|
||||
service_tier,
|
||||
)
|
||||
if service_tier
|
||||
else f"cache_read_input_token_cost_above_{threshold_str}_tokens"
|
||||
)
|
||||
|
||||
if cache_creation_tiered_key in model_info:
|
||||
cache_creation_cost = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_tiered_key,
|
||||
cache_creation_cost,
|
||||
),
|
||||
)
|
||||
cache_creation_cost = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_tiered_key,
|
||||
cache_creation_cost,
|
||||
),
|
||||
)
|
||||
|
||||
if cache_creation_1hr_tiered_key in model_info:
|
||||
cache_creation_cost_above_1hr = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_1hr_tiered_key,
|
||||
cache_creation_cost_above_1hr,
|
||||
),
|
||||
)
|
||||
cache_creation_cost_above_1hr = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info,
|
||||
cache_creation_1hr_tiered_key,
|
||||
cache_creation_cost_above_1hr,
|
||||
),
|
||||
)
|
||||
|
||||
if cache_read_tiered_key in model_info:
|
||||
cache_read_cost = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info, cache_read_tiered_key, cache_read_cost
|
||||
),
|
||||
)
|
||||
cache_read_cost = cast(
|
||||
float,
|
||||
_get_cost_per_unit(
|
||||
model_info, cache_read_tiered_key, cache_read_cost
|
||||
),
|
||||
)
|
||||
|
||||
break
|
||||
except (IndexError, ValueError):
|
||||
|
|
|
|||
|
|
@ -744,6 +744,17 @@ def _count_content_list(
|
|||
thinking_text = str(c.get("thinking", ""))
|
||||
if thinking_text:
|
||||
num_tokens += count_function(thinking_text)
|
||||
elif c["type"] == "tool_reference":
|
||||
# Anthropic tool-search reference block: a lightweight pointer to
|
||||
# a deferred tool, e.g. {"type": "tool_reference", "tool_name": ...}.
|
||||
# The full tool definition is counted via the `tools` param, so we
|
||||
# only count the referenced name here. Without this branch,
|
||||
# token_counter raises on tool-search traffic; on the streaming
|
||||
# anthropic_messages path that nulls response_cost and causes the
|
||||
# proxy to drop the SpendLogs row entirely (silent cost undercount).
|
||||
tool_name = str(c.get("tool_name") or "")
|
||||
if tool_name:
|
||||
num_tokens += count_function(tool_name)
|
||||
else:
|
||||
content_type = (
|
||||
c.get("type", type(c).__name__)
|
||||
|
|
@ -752,7 +763,7 @@ def _count_content_list(
|
|||
)
|
||||
raise ValueError(
|
||||
f"Invalid content item type: {content_type}. "
|
||||
f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking)."
|
||||
f"Expected str or dict with 'type' field (text, image_url, tool_use, tool_result, thinking, tool_reference)."
|
||||
)
|
||||
return num_tokens
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions`
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import (
|
||||
Any,
|
||||
Coroutine,
|
||||
|
|
@ -22,7 +23,9 @@ from litellm.litellm_core_utils.prompt_templates.factory import _parse_mime_type
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionFileObject,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
ChatCompletionVideoObject,
|
||||
ChatCompletionVideoUrlObject,
|
||||
)
|
||||
|
|
@ -101,26 +104,18 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
|
|||
) -> dict:
|
||||
_tools = non_default_params.pop("tools", None)
|
||||
if _tools is not None:
|
||||
# remove 'additionalProperties' from tools
|
||||
_tools = _remove_additional_properties(_tools)
|
||||
# remove 'strict' from tools
|
||||
_tools = _remove_strict_from_schema(_tools)
|
||||
if isinstance(_tools, list):
|
||||
_tools = self._convert_custom_tools_to_function_tools(_tools)
|
||||
if _tools is not None:
|
||||
non_default_params["tools"] = _tools
|
||||
|
||||
# Handle thinking parameter - convert Anthropic-style to OpenAI-style reasoning_effort
|
||||
# vLLM is OpenAI-compatible, so it understands reasoning_effort, not thinking
|
||||
# Reference: https://github.com/BerriAI/litellm/issues/19761
|
||||
thinking = non_default_params.pop("thinking", None)
|
||||
if thinking is not None and isinstance(thinking, dict):
|
||||
if thinking.get("type") == "enabled":
|
||||
# Only convert if reasoning_effort not already set
|
||||
if "reasoning_effort" not in non_default_params:
|
||||
budget_tokens = thinking.get("budget_tokens", 0)
|
||||
# Map budget_tokens to reasoning_effort level
|
||||
# Same logic as Anthropic adapter (translate_anthropic_thinking_to_reasoning_effort)
|
||||
if budget_tokens >= 10000:
|
||||
non_default_params["reasoning_effort"] = "high"
|
||||
elif budget_tokens >= 5000:
|
||||
|
|
@ -137,20 +132,13 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
|
|||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: Optional[str], api_key: Optional[str]
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
api_base = api_base or get_secret_str("HOSTED_VLLM_API_BASE") # type: ignore
|
||||
api_base = api_base or get_secret_str("HOSTED_VLLM_API_BASE")
|
||||
dynamic_api_key = (
|
||||
api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key"
|
||||
) # vllm does not require an api key
|
||||
)
|
||||
return api_base, dynamic_api_key
|
||||
|
||||
def _is_video_file(self, content_item: ChatCompletionFileObject) -> bool:
|
||||
"""
|
||||
Check if the file is a video
|
||||
|
||||
- format: video/<extension>
|
||||
- file_data: base64 encoded video data
|
||||
- file_id: infer mp4 from extension
|
||||
"""
|
||||
file = content_item.get("file", {})
|
||||
format = file.get("format")
|
||||
file_data = file.get("file_data")
|
||||
|
|
@ -205,29 +193,82 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
|
|||
"""
|
||||
Support translating:
|
||||
- video files from file_id or file_data to video_url
|
||||
- thinking_blocks on assistant messages to content blocks
|
||||
- thinking_blocks on assistant messages are removed, and content lists
|
||||
are converted to strings for vLLM compatibility
|
||||
"""
|
||||
for message in messages:
|
||||
if message["role"] == "assistant":
|
||||
thinking_blocks = message.pop("thinking_blocks", None) # type: ignore
|
||||
if thinking_blocks:
|
||||
new_content: list = [
|
||||
(
|
||||
{
|
||||
"type": block["type"],
|
||||
"thinking": block.get("thinking", ""),
|
||||
message.pop("thinking_blocks", None)
|
||||
existing_content = message.get("content")
|
||||
if isinstance(existing_content, list):
|
||||
text_parts = []
|
||||
tool_calls: list[ChatCompletionAssistantToolCall] = []
|
||||
content_blocks: list[object] = []
|
||||
has_structured_content = False
|
||||
for c in existing_content: # any-ok: untyped content
|
||||
if (
|
||||
isinstance(c, dict) # any-ok: untyped content
|
||||
and c.get("type") == "text" # any-ok: untyped content
|
||||
):
|
||||
text_parts.append( # any-ok: untyped content
|
||||
c.get("text", "") # any-ok: untyped content
|
||||
)
|
||||
content_blocks.append(c) # any-ok: untyped content
|
||||
elif (
|
||||
isinstance(c, dict) # any-ok: untyped content
|
||||
and c.get("type") == "tool_use" # any-ok: untyped content
|
||||
):
|
||||
tool_input = c.get("input", {}) # any-ok: untyped content
|
||||
tool_calls.append(
|
||||
ChatCompletionAssistantToolCall(
|
||||
id=c.get("id"), # any-ok: untyped content
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=c.get("name"), # any-ok: untyped content
|
||||
arguments=(
|
||||
tool_input
|
||||
if isinstance(
|
||||
tool_input, # any-ok: untyped content
|
||||
str, # any-ok: untyped content
|
||||
)
|
||||
else json.dumps(
|
||||
tool_input # any-ok: untyped content
|
||||
)
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
else:
|
||||
content_blocks.append(c) # any-ok: untyped content
|
||||
has_structured_content = True
|
||||
if tool_calls:
|
||||
existing_tool_calls = message.get("tool_calls")
|
||||
if isinstance(existing_tool_calls, list):
|
||||
existing_tool_call_ids = {
|
||||
tool_call.get("id") # any-ok: untyped content
|
||||
for tool_call in existing_tool_calls
|
||||
if isinstance(
|
||||
tool_call, dict
|
||||
) # any-ok: untyped content
|
||||
and tool_call.get("id")
|
||||
is not None # any-ok: untyped content
|
||||
}
|
||||
if block.get("type") == "thinking"
|
||||
else {"type": block["type"], "data": block.get("data", "")}
|
||||
)
|
||||
for block in thinking_blocks
|
||||
]
|
||||
existing_content = message.get("content")
|
||||
if isinstance(existing_content, str):
|
||||
new_content.append({"type": "text", "text": existing_content})
|
||||
elif isinstance(existing_content, list):
|
||||
new_content.extend(existing_content)
|
||||
message["content"] = new_content # type: ignore
|
||||
new_tool_calls = [
|
||||
tool_call
|
||||
for tool_call in tool_calls
|
||||
if tool_call.get("id") not in existing_tool_call_ids
|
||||
]
|
||||
if new_tool_calls:
|
||||
message["tool_calls"] = (
|
||||
existing_tool_calls + new_tool_calls
|
||||
)
|
||||
else:
|
||||
message["tool_calls"] = tool_calls
|
||||
content_str = "\n".join(text_parts) # any-ok: untyped content
|
||||
new_content = (
|
||||
content_blocks if has_structured_content else content_str
|
||||
)
|
||||
message["content"] = new_content # type: ignore[typeddict-item]
|
||||
elif message["role"] == "user":
|
||||
message_content = message.get("content")
|
||||
if message_content and isinstance(message_content, list):
|
||||
|
|
@ -243,6 +284,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
|
|||
message_content[idx] = self._convert_file_to_video_url(
|
||||
content_item
|
||||
)
|
||||
|
||||
if is_async:
|
||||
return super()._transform_messages(
|
||||
messages, model, is_async=cast(Literal[True], True)
|
||||
|
|
|
|||
|
|
@ -50,11 +50,15 @@ class OpenrouterConfig(OpenAIGPTConfig):
|
|||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
non_default_params: dict[str, object],
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
# OpenRouter expects "xhigh" instead of "max" for reasoning_effort.
|
||||
if non_default_params.get("reasoning_effort") == "max":
|
||||
non_default_params = {**non_default_params, "reasoning_effort": "xhigh"}
|
||||
|
||||
mapped_openai_params = super().map_openai_params(
|
||||
non_default_params, optional_params, model, drop_params
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2844,6 +2844,7 @@ async def make_call(
|
|||
sync_stream=False,
|
||||
logging_obj=logging_obj,
|
||||
response_headers=response.headers,
|
||||
response=response, # any-ok: untyped stream
|
||||
)
|
||||
# LOGGING
|
||||
logging_obj.post_call(
|
||||
|
|
@ -2887,6 +2888,7 @@ def make_sync_call(
|
|||
sync_stream=True,
|
||||
logging_obj=logging_obj,
|
||||
response_headers=response.headers,
|
||||
response=response, # any-ok: untyped stream
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -3348,12 +3350,14 @@ class ModelResponseIterator:
|
|||
sync_stream: bool,
|
||||
logging_obj: LoggingClass,
|
||||
response_headers: Optional[Dict[str, str]] = None,
|
||||
response: httpx.Response | None = None,
|
||||
):
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
check_is_function_call,
|
||||
)
|
||||
|
||||
self.streaming_response = streaming_response
|
||||
self.response = response
|
||||
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
|
||||
self.accumulated_json = ""
|
||||
self.sent_first_chunk = False
|
||||
|
|
@ -3655,3 +3659,47 @@ class ModelResponseIterator:
|
|||
raise StopAsyncIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
|
||||
|
||||
async def aclose(self) -> None:
|
||||
iterator = getattr( # any-ok: untyped stream
|
||||
self,
|
||||
"async_response_iterator",
|
||||
self.streaming_response, # any-ok: untyped stream
|
||||
)
|
||||
if iterator is not None and hasattr( # any-ok: untyped stream
|
||||
iterator, "aclose" # any-ok: untyped stream
|
||||
):
|
||||
try:
|
||||
await iterator.aclose() # any-ok: untyped stream
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.debug(
|
||||
"ModelResponseIterator.aclose: error closing iterator: %s", e
|
||||
)
|
||||
if self.response is not None:
|
||||
try:
|
||||
await self.response.aclose()
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.debug(
|
||||
"ModelResponseIterator.aclose: error closing response: %s", e
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
iterator = getattr( # any-ok: untyped stream
|
||||
self, "response_iterator", self.streaming_response # any-ok: untyped stream
|
||||
)
|
||||
if iterator is not None and hasattr( # any-ok: untyped stream
|
||||
iterator, "close" # any-ok: untyped stream
|
||||
):
|
||||
try:
|
||||
iterator.close() # any-ok: untyped stream
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.debug(
|
||||
"ModelResponseIterator.close: error closing iterator: %s", e
|
||||
)
|
||||
if self.response is not None:
|
||||
try:
|
||||
self.response.close()
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.debug(
|
||||
"ModelResponseIterator.close: error closing response: %s", e
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2528,6 +2528,100 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"azure_ai/gpt-5.5": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"source": "https://ai.azure.com/catalog/models/gpt-5.5",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": 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_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false
|
||||
},
|
||||
"azure_ai/gpt-5.5-2026-04-23": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"source": "https://ai.azure.com/catalog/models/gpt-5.5",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": 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_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false
|
||||
},
|
||||
"azure_ai/gpt-5.4": {
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 5e-07,
|
||||
|
|
@ -10068,6 +10162,8 @@
|
|||
},
|
||||
"claude-sonnet-4-5": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
|
|
@ -10097,6 +10193,8 @@
|
|||
},
|
||||
"claude-sonnet-4-5-20250929": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
|
|
@ -10127,6 +10225,7 @@
|
|||
},
|
||||
"claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "anthropic",
|
||||
|
|
@ -10155,6 +10254,8 @@
|
|||
},
|
||||
"claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
|
|
@ -25103,6 +25204,21 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"mistral/mistral-medium-3-5": {
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "mistral",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"mistral/mistral-small": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "mistral",
|
||||
|
|
@ -42456,4 +42572,105 @@
|
|||
"supports_reasoning": true,
|
||||
"source": "https://serverless.tensormesh.ai/v1/models/openrouter"
|
||||
}
|
||||
}
|
||||
,
|
||||
"deepseek-v4-flash": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"input_cost_per_token_cache_hit": 2.8e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"deepseek-v4-pro": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 3.625e-09,
|
||||
"input_cost_per_token": 4.35e-07,
|
||||
"input_cost_per_token_cache_hit": 3.625e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.7e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"deepseek/deepseek-v4-flash": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"input_cost_per_token_cache_hit": 2.8e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"deepseek/deepseek-v4-pro": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 3.625e-09,
|
||||
"input_cost_per_token": 4.35e-07,
|
||||
"input_cost_per_token_cache_hit": 3.625e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.7e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1421,8 +1421,11 @@ class MCPServerManager:
|
|||
"No allowed MCP Servers found for user api key auth."
|
||||
)
|
||||
return list(combined_servers)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}.")
|
||||
except Exception: # noqa: BLE001
|
||||
verbose_logger.exception(
|
||||
"Failed to get allowed MCP servers; team-level object_permission "
|
||||
"grants may be dropped. Falling back to global servers only."
|
||||
)
|
||||
return allow_all_server_ids
|
||||
|
||||
async def resolve_toolset_tool_permissions(
|
||||
|
|
|
|||
|
|
@ -1036,7 +1036,14 @@ if MCP_AVAILABLE:
|
|||
allowed_mcp_servers: List[MCPServer],
|
||||
) -> List[MCPServer]:
|
||||
"""
|
||||
Get the filtered MCP servers from the MCP server names
|
||||
Get the filtered MCP servers from the MCP server names.
|
||||
|
||||
Fails closed when ``mcp_servers`` is explicitly provided (path- or
|
||||
header-derived) but none of the names resolve to a server alias or
|
||||
access group the caller can access. The previous behavior returned
|
||||
the full ``allowed_mcp_servers`` set, which silently widened scope
|
||||
when a client targeted ``/mcp/<unknown>/`` and made URL/header
|
||||
namespacing appear to work when it did not.
|
||||
"""
|
||||
|
||||
filtered_server: dict[str, MCPServer] = {}
|
||||
|
|
@ -1076,6 +1083,17 @@ if MCP_AVAILABLE:
|
|||
if filtered_server:
|
||||
return list(filtered_server.values())
|
||||
|
||||
if mcp_servers is not None:
|
||||
# Caller asked for a specific scope but nothing resolved. Fail
|
||||
# closed so URL/header namespacing cannot silently fall back to
|
||||
# the caller's full allowed-server set.
|
||||
verbose_logger.debug(
|
||||
"MCP scope filter resolved to no servers for requested names %s; "
|
||||
"returning empty list (fail-closed).",
|
||||
mcp_servers,
|
||||
)
|
||||
return []
|
||||
|
||||
return allowed_mcp_servers
|
||||
|
||||
def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.repositories.object_permission_repository import ObjectPermissionRe
|
|||
from litellm.router import Router
|
||||
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
|
||||
from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import get_valid_models
|
||||
|
||||
_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)
|
||||
|
|
@ -308,10 +309,21 @@ def get_known_models_from_wildcard(
|
|||
# add model prefix to wildcard models
|
||||
wildcard_models = [f"{model_prefix}{model}" for model in wildcard_models]
|
||||
|
||||
known_providers = {provider.value for provider in LlmProviders}
|
||||
suffix_appended_wildcard_models = []
|
||||
for model in wildcard_models:
|
||||
if not model.startswith(wildcard_provider_prefix):
|
||||
model = f"{wildcard_provider_prefix}/{model}"
|
||||
# `get_provider_models` returns provider-prefixed ids (e.g. "ollama/gemma3:1b").
|
||||
# When the wildcard uses a custom prefix (e.g. "ollama_server1/*" to distinguish
|
||||
# multiple instances), replace that existing provider prefix instead of stacking
|
||||
# both, which would otherwise yield an uncallable "ollama_server1/ollama/gemma3:1b".
|
||||
# Only strip the leading segment when it is a known provider, so ids whose first
|
||||
# segment is an org rather than a provider (e.g. "meta-llama/Llama-3-8B") keep it.
|
||||
leading, sep, model_suffix = model.partition("/")
|
||||
if sep and leading in known_providers:
|
||||
model = f"{wildcard_provider_prefix}/{model_suffix}"
|
||||
else:
|
||||
model = f"{wildcard_provider_prefix}/{model}"
|
||||
suffix_appended_wildcard_models.append(model)
|
||||
return suffix_appended_wildcard_models or []
|
||||
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.constants import (
|
|||
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
|
||||
DEFAULT_MAX_RECURSE_DEPTH,
|
||||
LITELLM_DETAILED_TIMING,
|
||||
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
|
||||
MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
|
||||
STREAM_SSE_DATA_PREFIX,
|
||||
)
|
||||
|
|
@ -67,7 +68,12 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
ProxyConfig = Any
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
Usage,
|
||||
)
|
||||
|
||||
# Datadog streaming spans are a no-op when ddtrace is not enabled, but the
|
||||
# ``with tracer.trace(...)`` context manager still allocates a NullSpan and
|
||||
|
|
@ -77,6 +83,77 @@ from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
|
|||
_DD_STREAMING_TRACE_ENABLED = not isinstance(tracer, NullTracer)
|
||||
|
||||
|
||||
_CLIENT_DISCONNECTED_ERROR_INFORMATION: StandardLoggingPayloadErrorInformation = {
|
||||
"error_code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED),
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
}
|
||||
|
||||
|
||||
def _apply_client_disconnect_metadata(target_metadata: dict[str, object]) -> None:
|
||||
target_metadata["client_disconnected"] = True
|
||||
target_metadata["error_information"] = dict(_CLIENT_DISCONNECTED_ERROR_INFORMATION)
|
||||
|
||||
|
||||
async def _record_streaming_client_disconnect_if_needed(
|
||||
request: Request | None,
|
||||
request_data: dict,
|
||||
client_disconnected: bool = False,
|
||||
) -> bool:
|
||||
if not client_disconnected:
|
||||
if request is None:
|
||||
return False
|
||||
try:
|
||||
disconnected = await request.is_disconnected()
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
if not disconnected:
|
||||
return False
|
||||
|
||||
logging_obj = request_data.get("litellm_logging_obj") # any-ok: untyped request
|
||||
if logging_obj is not None: # any-ok: untyped request
|
||||
litellm_params = (
|
||||
logging_obj.model_call_details.setdefault( # any-ok: untyped request
|
||||
"litellm_params", {}
|
||||
)
|
||||
)
|
||||
_apply_client_disconnect_metadata(
|
||||
litellm_params.setdefault("metadata", {}) # any-ok: untyped request
|
||||
)
|
||||
_apply_client_disconnect_metadata(
|
||||
logging_obj.model_call_details.setdefault( # any-ok: untyped request
|
||||
"metadata", {}
|
||||
)
|
||||
)
|
||||
|
||||
_apply_client_disconnect_metadata(
|
||||
request_data.setdefault("metadata", {}) # any-ok: untyped request
|
||||
)
|
||||
litellm_params = request_data.setdefault( # any-ok: untyped request
|
||||
"litellm_params", {} # any-ok: untyped request
|
||||
)
|
||||
_apply_client_disconnect_metadata(
|
||||
litellm_params.setdefault("metadata", {}) # any-ok: untyped request
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Recorded streaming client disconnect with error_code=499 for litellm_call_id=%s",
|
||||
request_data.get("litellm_call_id"), # any-ok: untyped request
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None:
|
||||
pending_tasks = [task for task in tasks if not task.done()] # any-ok: untyped task
|
||||
for task in pending_tasks: # any-ok: untyped task
|
||||
task.cancel() # any-ok: untyped task
|
||||
for task in pending_tasks: # any-ok: untyped task
|
||||
try:
|
||||
await task # any-ok: untyped request
|
||||
except (asyncio.CancelledError, Exception): # noqa: BLE001
|
||||
pass
|
||||
|
||||
|
||||
def _serialize_http_exception_detail(
|
||||
detail: Any,
|
||||
) -> Tuple[str, Optional[dict]]:
|
||||
|
|
@ -242,20 +319,6 @@ def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict:
|
|||
return default_error
|
||||
|
||||
|
||||
async def _aclose_upstream_response(response: Any) -> None:
|
||||
"""Release the upstream HTTP connection when a stream ends for any
|
||||
reason, including client disconnect. Mirrors the finally block of
|
||||
async_data_generator in proxy_server.py."""
|
||||
with anyio.CancelScope(shield=True):
|
||||
if hasattr(response, "aclose"):
|
||||
try:
|
||||
await response.aclose()
|
||||
except BaseException as e:
|
||||
verbose_proxy_logger.debug(
|
||||
"error closing upstream response stream: %s", e
|
||||
)
|
||||
|
||||
|
||||
class _UpstreamClosingStreamingResponse(StreamingResponse):
|
||||
"""StreamingResponse that always closes its body iterator and the wrapped
|
||||
upstream generator.
|
||||
|
|
@ -1338,19 +1401,24 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_model=user_model,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
tasks.append(llm_call)
|
||||
llm_call_task = asyncio.create_task(llm_call) # any-ok: untyped task
|
||||
tasks.append(llm_call_task) # any-ok: untyped task
|
||||
|
||||
# wait for call to end
|
||||
llm_responses = asyncio.gather(
|
||||
*tasks
|
||||
) # run the moderation check in parallel to the actual llm api call
|
||||
|
||||
if general_settings.get("cancel_on_disconnect", False):
|
||||
responses = await _await_llm_call_cancelling_on_disconnect(
|
||||
request, llm_responses
|
||||
)
|
||||
else:
|
||||
responses = await llm_responses
|
||||
try:
|
||||
if general_settings.get( # any-ok: untyped request
|
||||
"cancel_on_disconnect", False
|
||||
):
|
||||
responses = await _await_llm_call_cancelling_on_disconnect( # any-ok: untyped request
|
||||
request, llm_responses # any-ok: untyped task
|
||||
)
|
||||
else:
|
||||
responses = await llm_responses # any-ok: untyped request
|
||||
finally:
|
||||
await _cancel_pending_gather_tasks(tasks) # any-ok: untyped task
|
||||
|
||||
response = responses[1]
|
||||
|
||||
|
|
@ -1526,6 +1594,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=self.data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
return await create_response(
|
||||
|
|
@ -1539,6 +1608,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=self.data,
|
||||
request=request,
|
||||
)
|
||||
if route_type == "aresponses":
|
||||
# Streaming /v1/responses returns here without
|
||||
|
|
@ -2383,6 +2453,41 @@ class ProxyBaseLLMRequestProcessing:
|
|||
else:
|
||||
return chunk
|
||||
|
||||
@staticmethod
|
||||
async def _finalize_streaming_generator_cleanup(
|
||||
request: Request | None,
|
||||
request_data: dict,
|
||||
response: Any,
|
||||
stream_completed: bool = False,
|
||||
client_disconnected: bool = False,
|
||||
) -> None:
|
||||
with anyio.CancelScope(shield=True):
|
||||
should_record_client_disconnect = client_disconnected or (
|
||||
not stream_completed
|
||||
)
|
||||
recorded_client_disconnect = False
|
||||
if should_record_client_disconnect:
|
||||
recorded_client_disconnect = (
|
||||
await _record_streaming_client_disconnect_if_needed(
|
||||
request,
|
||||
request_data, # any-ok: untyped request
|
||||
client_disconnected, # any-ok: untyped request
|
||||
)
|
||||
)
|
||||
if recorded_client_disconnect:
|
||||
ProxyLogging._fire_deferred_stream_logging(
|
||||
request_data # any-ok: untyped request
|
||||
)
|
||||
|
||||
if hasattr(response, "aclose"): # any-ok: untyped request
|
||||
try:
|
||||
await response.aclose() # any-ok: untyped request
|
||||
except BaseException as e: # noqa: BLE001
|
||||
verbose_proxy_logger.debug(
|
||||
"async_streaming_data_generator: error closing response stream: %s",
|
||||
e,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def async_streaming_data_generator(
|
||||
response: Any,
|
||||
|
|
@ -2392,6 +2497,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
*,
|
||||
serialize_chunk: StreamChunkSerializer,
|
||||
serialize_error: StreamErrorSerializer,
|
||||
request: Request | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
|
||||
|
|
@ -2416,6 +2522,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
and not cost_injection_enabled
|
||||
)
|
||||
debug_enabled = verbose_proxy_logger.isEnabledFor(logging.DEBUG)
|
||||
stream_completed = False
|
||||
client_disconnected = False
|
||||
try:
|
||||
str_so_far = ""
|
||||
async for (
|
||||
|
|
@ -2463,6 +2571,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
)
|
||||
)
|
||||
yield serialize_chunk(chunk)
|
||||
stream_completed = True
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
# Client disconnected mid-stream. CancelledError / GeneratorExit
|
||||
# are BaseException and bypass the success/failure logging
|
||||
|
|
@ -2470,9 +2579,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# release it here. This is the outermost generator Starlette closes
|
||||
# on disconnect, so the nested iterator hook (which only sees
|
||||
# GeneratorExit on GC) cannot own the refund.
|
||||
proxy_logging_obj._release_max_parallel_requests_on_disconnect(
|
||||
user_api_key_dict
|
||||
)
|
||||
if not stream_completed:
|
||||
proxy_logging_obj._release_max_parallel_requests_on_disconnect(
|
||||
user_api_key_dict
|
||||
)
|
||||
client_disconnected = True
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
@ -2501,9 +2612,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", 500),
|
||||
)
|
||||
stream_completed = True
|
||||
yield serialize_error(proxy_exception)
|
||||
finally:
|
||||
await _aclose_upstream_response(response)
|
||||
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
|
||||
request=request,
|
||||
request_data=request_data, # any-ok: untyped request
|
||||
response=response, # any-ok: untyped request
|
||||
stream_completed=stream_completed,
|
||||
client_disconnected=client_disconnected,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def async_sse_data_generator(
|
||||
|
|
@ -2511,6 +2629,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_data: dict,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request: Request | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Anthropic /messages and Google /generateContent streaming data generator require SSE events.
|
||||
|
|
@ -2529,6 +2648,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
serialize_error=lambda proxy_exc: (
|
||||
f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n"
|
||||
),
|
||||
request=request,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -25,6 +25,16 @@ async def get_ui_config():
|
|||
or general_settings.get("auto_redirect_ui_login_to_sso", False) is True
|
||||
)
|
||||
admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"
|
||||
hide_default_credentials_hint = bool( # any-ok: untyped settings
|
||||
os.getenv( # any-ok: untyped settings
|
||||
"LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false"
|
||||
).lower()
|
||||
== "true"
|
||||
or general_settings.get( # any-ok: untyped settings
|
||||
"hide_default_credentials_hint", False
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
sso_configured = _has_user_setup_sso()
|
||||
|
||||
|
|
@ -38,6 +48,7 @@ async def get_ui_config():
|
|||
auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso,
|
||||
admin_ui_disabled=admin_ui_disabled,
|
||||
sso_configured=sso_configured,
|
||||
hide_default_credentials_hint=hide_default_credentials_hint, # any-ok: untyped settings
|
||||
is_control_plane=is_control_plane,
|
||||
workers=proxy_config.worker_registry if is_control_plane else [],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -107,6 +107,7 @@ async def google_stream_generate_content(
|
|||
data["stream"] = True
|
||||
# google-genai SDK (?alt=sse) must not receive OpenAI's data: [DONE] terminator.
|
||||
data["_litellm_skip_openai_stream_done"] = True
|
||||
data["_litellm_raw_sse_stream"] = True # any-ok: untyped request
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -741,6 +741,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
|
||||
For multiple messages in /chat/completions, we'll need to call them in parallel.
|
||||
"""
|
||||
# Respect the configured event hook. In `logging_only` mode (and any config that
|
||||
# excludes pre_call) the live request must not be masked - masking is applied to a
|
||||
# copy at logging time via `async_logging_hook`. Without this gate the request sent
|
||||
# to the model would carry anonymization tokens and the response would echo them.
|
||||
if (
|
||||
self.should_run_guardrail(
|
||||
data=data, # any-ok: untyped request
|
||||
event_type=GuardrailEventHooks.pre_call, # any-ok: untyped request
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return data # any-ok: untyped request
|
||||
|
||||
try:
|
||||
content_safety = data.get("content_safety", None)
|
||||
|
|
|
|||
|
|
@ -162,6 +162,8 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = (
|
|||
"secret_fields",
|
||||
"_guardrail_pipelines",
|
||||
"_pipeline_managed_guardrails",
|
||||
"client_disconnected",
|
||||
"error_information",
|
||||
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
|
|
@ -390,38 +390,24 @@ def _adjust_dates_for_timezone(
|
|||
timezone_offset_minutes: Optional[int],
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Adjust date range to account for timezone differences.
|
||||
Pass-through for the local date range; the timezone offset is intentionally ignored here.
|
||||
|
||||
The database stores dates in UTC. When a user in a different timezone
|
||||
selects a local date range, we need to expand the UTC query range to
|
||||
capture all records that fall within their local date range.
|
||||
The aggregation table (e.g. LiteLLM_DailyUserSpend) stores spend in whole-UTC-day
|
||||
buckets keyed on date as YYYY-MM-DD. Any conversion from a local date range to a
|
||||
UTC date range using only date arithmetic must round to whole UTC days, allowing up
|
||||
to 24h of slop at each boundary. The previous implementation expanded the SQL range
|
||||
by an extra full UTC day on whichever side the offset pointed, which pulled in 24h
|
||||
of unrelated bucket data per boundary and produced approximately 100% over-counting
|
||||
on single-day queries (e.g. IST May 29 returning UTC May 28 + UTC May 29 in full).
|
||||
Sums of single-day queries then exceeded the equivalent multi-day aggregate, which
|
||||
is mathematically impossible.
|
||||
|
||||
Args:
|
||||
start_date: Start date in YYYY-MM-DD format (user's local date)
|
||||
end_date: End date in YYYY-MM-DD format (user's local date)
|
||||
timezone_offset_minutes: Minutes behind UTC (positive = west of UTC)
|
||||
This matches JavaScript's Date.getTimezoneOffset() convention.
|
||||
For example: PST = +480 (8 hours * 60 = 480 minutes behind UTC)
|
||||
|
||||
Returns:
|
||||
Tuple of (adjusted_start_date, adjusted_end_date) in YYYY-MM-DD format
|
||||
Treating the local date as the UTC date trades a small one-time boundary slop for
|
||||
correct, monotonic, additive results across single-day and multi-day queries. A
|
||||
later fix can introduce hour-level buckets or pro-rata weighting on adjacent UTC
|
||||
days; both require data the current schema does not store.
|
||||
"""
|
||||
if timezone_offset_minutes is None or timezone_offset_minutes == 0:
|
||||
return start_date, end_date
|
||||
|
||||
start = datetime.strptime(start_date, "%Y-%m-%d")
|
||||
end = datetime.strptime(end_date, "%Y-%m-%d")
|
||||
|
||||
if timezone_offset_minutes > 0:
|
||||
# West of UTC (Americas): local evening extends into next UTC day
|
||||
# e.g., Feb 4 23:59 PST = Feb 5 07:59 UTC
|
||||
end = end + timedelta(days=1)
|
||||
else:
|
||||
# East of UTC (Asia/Europe): local morning starts in previous UTC day
|
||||
# e.g., Feb 4 00:00 IST = Feb 3 18:30 UTC
|
||||
start = start - timedelta(days=1)
|
||||
|
||||
return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")
|
||||
return start_date, end_date
|
||||
|
||||
|
||||
def _build_where_conditions(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import os
|
|||
import re
|
||||
import secrets
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, cast
|
||||
|
||||
|
|
@ -59,6 +60,9 @@ from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_k
|
|||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_check_passthrough_routes_caller_permission,
|
||||
_is_user_org_admin_for_team,
|
||||
|
|
@ -3225,6 +3229,69 @@ async def delete_key_fn(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
async def _get_model_max_budget_current_spend(
|
||||
api_key_hash: str,
|
||||
model: str,
|
||||
budget_config: BudgetConfig,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> float:
|
||||
virtual_key_model_spend_cache_key = (
|
||||
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:"
|
||||
f"{api_key_hash}:{model}:{budget_config.budget_duration}"
|
||||
)
|
||||
current_spend: float | None = (
|
||||
await user_api_key_cache.async_get_cache( # any-ok: untyped dump
|
||||
key=virtual_key_model_spend_cache_key,
|
||||
)
|
||||
)
|
||||
if current_spend is None:
|
||||
model_without_prefix = model.split("/")[-1] if "/" in model else model
|
||||
virtual_key_model_spend_cache_key = (
|
||||
f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:"
|
||||
f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}"
|
||||
)
|
||||
current_spend = (
|
||||
await user_api_key_cache.async_get_cache( # any-ok: untyped dump
|
||||
key=virtual_key_model_spend_cache_key,
|
||||
)
|
||||
)
|
||||
try:
|
||||
return float(current_spend or 0.0) # any-ok: untyped dump
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
async def _build_model_max_budget_usage(
|
||||
api_key_hash: str,
|
||||
model_max_budget: Mapping[str, Mapping[str, object]],
|
||||
user_api_key_cache: UserApiKeyCache | None,
|
||||
) -> dict[str, dict[str, object]]:
|
||||
if user_api_key_cache is None or not model_max_budget:
|
||||
return {}
|
||||
|
||||
result: dict[str, dict[str, object]] = {}
|
||||
for model, budget_info in model_max_budget.items():
|
||||
try:
|
||||
budget_config = BudgetConfig.model_validate(budget_info)
|
||||
if budget_config.budget_duration is None:
|
||||
continue
|
||||
duration_in_seconds(budget_config.budget_duration)
|
||||
except Exception: # noqa: BLE001
|
||||
continue
|
||||
spend = await _get_model_max_budget_current_spend(
|
||||
api_key_hash=api_key_hash,
|
||||
model=model,
|
||||
budget_config=budget_config,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
result[model] = {
|
||||
"current_spend": round(spend, 4),
|
||||
"budget_limit": budget_config.max_budget,
|
||||
"time_period": budget_config.budget_duration,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v2/key/info",
|
||||
tags=["key management"],
|
||||
|
|
@ -3252,7 +3319,7 @@ async def info_key_fn_v2(
|
|||
-d {"keys": ["sk-1", "sk-2", "sk-3"]}
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
if prisma_client is None:
|
||||
|
|
@ -3298,7 +3365,29 @@ async def info_key_fn_v2(
|
|||
k_dict = k.model_dump()
|
||||
except Exception:
|
||||
k_dict = k.dict()
|
||||
k_dict.pop("token", None)
|
||||
k_token_hash = k_dict.pop("token", None) # any-ok: untyped dump
|
||||
|
||||
model_max_budget = (
|
||||
k_dict.get("model_max_budget") or {} # any-ok: untyped dump
|
||||
)
|
||||
budget_table = (
|
||||
k_dict.get("litellm_budget_table") or {} # any-ok: untyped dump
|
||||
)
|
||||
if not model_max_budget and isinstance( # any-ok: untyped dump
|
||||
budget_table, dict # any-ok: untyped dump
|
||||
):
|
||||
model_max_budget = (
|
||||
budget_table.get("model_max_budget") or {} # any-ok: untyped dump
|
||||
)
|
||||
if model_max_budget and k_token_hash: # any-ok: untyped dump
|
||||
k_dict["model_max_budget_usage"] = ( # any-ok: untyped dump
|
||||
await _build_model_max_budget_usage( # any-ok: untyped dump
|
||||
api_key_hash=k_token_hash, # any-ok: untyped dump
|
||||
model_max_budget=model_max_budget, # any-ok: untyped dump
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
)
|
||||
|
||||
filtered_key_info.append(k_dict)
|
||||
return {"key": data.keys, "info": filtered_key_info}
|
||||
|
||||
|
|
@ -3336,7 +3425,7 @@ async def info_key_fn(
|
|||
-H "Authorization: Bearer sk-test-example-key-123"
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
if prisma_client is None:
|
||||
|
|
@ -3381,7 +3470,28 @@ async def info_key_fn(
|
|||
except Exception:
|
||||
# if using pydantic v1
|
||||
key_info = key_info.dict()
|
||||
key_info.pop("token")
|
||||
key_token_hash = key_info.pop("token") # any-ok: untyped dump
|
||||
|
||||
model_max_budget = (
|
||||
key_info.get("model_max_budget") or {} # any-ok: untyped dump
|
||||
)
|
||||
budget_table = (
|
||||
key_info.get("litellm_budget_table") or {} # any-ok: untyped dump
|
||||
)
|
||||
if not model_max_budget and isinstance( # any-ok: untyped dump
|
||||
budget_table, dict # any-ok: untyped dump
|
||||
):
|
||||
model_max_budget = (
|
||||
budget_table.get("model_max_budget") or {} # any-ok: untyped dump
|
||||
)
|
||||
if model_max_budget and key_token_hash: # any-ok: untyped dump
|
||||
key_info["model_max_budget_usage"] = ( # any-ok: untyped dump
|
||||
await _build_model_max_budget_usage( # any-ok: untyped dump
|
||||
api_key_hash=key_token_hash, # any-ok: untyped dump
|
||||
model_max_budget=model_max_budget, # any-ok: untyped dump
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
)
|
||||
|
||||
# Attach object_permission if object_permission_id is set
|
||||
key_info = await attach_object_permission_to_dict(key_info, prisma_client)
|
||||
|
|
|
|||
|
|
@ -888,16 +888,26 @@ async def proxy_startup_event(app: FastAPI):
|
|||
|
||||
asyncio.create_task(_run_pw_migration())
|
||||
|
||||
## use_redis_transaction_buffer: fall back to a standalone Redis (REDIS_* env)
|
||||
## when the proxy cache backend is not Redis ##
|
||||
transaction_buffer_redis_cache = redis_usage_cache
|
||||
if transaction_buffer_redis_cache is None:
|
||||
transaction_buffer_redis_cache = (
|
||||
ProxyStartupEvent._get_transaction_buffer_redis_cache(
|
||||
general_settings=general_settings # any-ok: untyped stream
|
||||
)
|
||||
)
|
||||
|
||||
ProxyStartupEvent._initialize_startup_logging(
|
||||
llm_router=llm_router,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
redis_usage_cache=redis_usage_cache,
|
||||
redis_usage_cache=transaction_buffer_redis_cache,
|
||||
)
|
||||
|
||||
## Validate use_redis_transaction_buffer requires Redis cache ##
|
||||
ProxyStartupEvent._validate_redis_transaction_buffer_config(
|
||||
general_settings=general_settings,
|
||||
redis_usage_cache=redis_usage_cache,
|
||||
redis_usage_cache=transaction_buffer_redis_cache,
|
||||
)
|
||||
|
||||
## SEMANTIC TOOL FILTER ##
|
||||
|
|
@ -7022,10 +7032,33 @@ def _format_streaming_sse_chunk(chunk: Union[str, bytes]) -> Union[str, bytes]:
|
|||
return f"data: {chunk}\n\n"
|
||||
|
||||
|
||||
_SSE_FRAME_DELIMITERS = ("\r\n\r\n", "\n\n", "\r\r")
|
||||
_MAX_RAW_SSE_BUFFER_CHARS = 8 * 1024 * 1024
|
||||
|
||||
|
||||
def _pop_complete_sse_frame(buffer: str) -> tuple[str | None, str]:
|
||||
delimiter_positions = [
|
||||
(position, delimiter)
|
||||
for delimiter in _SSE_FRAME_DELIMITERS
|
||||
if (position := buffer.find(delimiter)) != -1
|
||||
]
|
||||
if not delimiter_positions:
|
||||
return None, buffer
|
||||
|
||||
position, delimiter = min(delimiter_positions, key=lambda item: item[0])
|
||||
frame_end = position + len(delimiter)
|
||||
return buffer[:frame_end], buffer[frame_end:]
|
||||
|
||||
|
||||
async def async_data_generator(
|
||||
response, user_api_key_dict: UserAPIKeyAuth, request_data: dict
|
||||
response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_data: dict,
|
||||
request: Request | None = None,
|
||||
):
|
||||
verbose_proxy_logger.debug("inside generator")
|
||||
stream_completed = False
|
||||
client_disconnected = False
|
||||
try:
|
||||
error_message: Optional[str] = None
|
||||
requested_model_from_client = _get_client_requested_model_for_streaming(
|
||||
|
|
@ -7047,6 +7080,10 @@ async def async_data_generator(
|
|||
# happened to ship a streaming-iterator override (the default).
|
||||
needs_iterator_wrap = proxy_logging_obj.needs_iterator_wrap()
|
||||
needs_per_chunk_hook = proxy_logging_obj.needs_per_chunk_streaming_hook()
|
||||
is_raw_sse_stream = bool(
|
||||
request_data.get("_litellm_raw_sse_stream") # any-ok: untyped stream
|
||||
)
|
||||
raw_sse_buffer = ""
|
||||
|
||||
if needs_iterator_wrap:
|
||||
stream_iterator = proxy_logging_obj.async_post_call_streaming_iterator_hook(
|
||||
|
|
@ -7077,14 +7114,38 @@ async def async_data_generator(
|
|||
if isinstance(chunk, BaseModel):
|
||||
chunk = _serialize_streaming_chunk(chunk)
|
||||
elif isinstance(chunk, bytes):
|
||||
# Some upstream streaming iterators (e.g. AsyncGoogleGenAIGenerateContentStreamingIterator
|
||||
# for /v1beta/.../streamGenerateContent) yield raw SSE bytes from Gemini.
|
||||
# Decode to str so the f-string below does not emit a Python b'...' literal,
|
||||
# and pass already-formatted SSE through unchanged to avoid double "data:" prefix.
|
||||
chunk = chunk.decode("utf-8", errors="replace")
|
||||
if chunk.startswith(("data:", "event:", ":")):
|
||||
yield chunk if chunk.endswith("\n\n") else chunk + "\n\n"
|
||||
if is_raw_sse_stream:
|
||||
raw_sse_buffer += chunk
|
||||
while True:
|
||||
frame, raw_sse_buffer = _pop_complete_sse_frame(raw_sse_buffer)
|
||||
if frame is None:
|
||||
break
|
||||
yield frame # any-ok: untyped stream
|
||||
if len(raw_sse_buffer) > _MAX_RAW_SSE_BUFFER_CHARS:
|
||||
raise ValueError(
|
||||
"Raw SSE stream exceeded maximum buffered size without a frame delimiter"
|
||||
)
|
||||
continue
|
||||
if chunk.startswith(("data:", "event:", ":")):
|
||||
yield ( # any-ok: untyped stream
|
||||
chunk
|
||||
if chunk.endswith(_SSE_FRAME_DELIMITERS)
|
||||
else chunk + "\n\n"
|
||||
)
|
||||
continue
|
||||
elif isinstance(chunk, str) and is_raw_sse_stream: # any-ok: untyped stream
|
||||
raw_sse_buffer += chunk
|
||||
while True:
|
||||
frame, raw_sse_buffer = _pop_complete_sse_frame(raw_sse_buffer)
|
||||
if frame is None:
|
||||
break
|
||||
yield frame # any-ok: untyped stream
|
||||
if len(raw_sse_buffer) > _MAX_RAW_SSE_BUFFER_CHARS:
|
||||
raise ValueError(
|
||||
"Raw SSE stream exceeded maximum buffered size without a frame delimiter"
|
||||
)
|
||||
continue
|
||||
elif isinstance(chunk, str) and chunk.startswith("data: "):
|
||||
error_message = chunk
|
||||
break
|
||||
|
|
@ -7094,12 +7155,20 @@ async def async_data_generator(
|
|||
except Exception as e:
|
||||
yield f"data: {str(e)}\n\n"
|
||||
|
||||
stream_completed = True
|
||||
if not needs_iterator_wrap:
|
||||
# The iterator-wrap path fires deferred logging itself; fire it
|
||||
# here for the no-wrap fast path so non-callback deployments
|
||||
# still flush their post-stream logging.
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
|
||||
if raw_sse_buffer:
|
||||
yield ( # any-ok: untyped stream
|
||||
raw_sse_buffer
|
||||
if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS)
|
||||
else raw_sse_buffer + "\n\n"
|
||||
)
|
||||
|
||||
if error_message is not None:
|
||||
yield error_message
|
||||
# OpenAI-compatible streams terminate with data: [DONE]; Google GenAI (?alt=sse) does not.
|
||||
|
|
@ -7113,9 +7182,11 @@ async def async_data_generator(
|
|||
# it here. This is the outermost generator Starlette closes on
|
||||
# disconnect, so it fires reliably regardless of needs_iterator_wrap
|
||||
# (a nested iterator hook would only see GeneratorExit on GC).
|
||||
proxy_logging_obj._release_max_parallel_requests_on_disconnect(
|
||||
user_api_key_dict
|
||||
)
|
||||
if not stream_completed:
|
||||
proxy_logging_obj._release_max_parallel_requests_on_disconnect(
|
||||
user_api_key_dict
|
||||
)
|
||||
client_disconnected = True
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
@ -7149,30 +7220,33 @@ async def async_data_generator(
|
|||
code=getattr(e, "status_code", 500),
|
||||
)
|
||||
error_returned = json.dumps({"error": proxy_exception.to_dict()})
|
||||
stream_completed = True
|
||||
yield f"data: {error_returned}\n\n"
|
||||
finally:
|
||||
# Close the response stream to release the underlying HTTP connection
|
||||
# back to the connection pool. This prevents pool exhaustion when
|
||||
# clients disconnect mid-stream.
|
||||
# Shield from cancellation so the close awaits can complete.
|
||||
with anyio.CancelScope(shield=True):
|
||||
if hasattr(response, "aclose"):
|
||||
try:
|
||||
await response.aclose()
|
||||
except BaseException as e:
|
||||
verbose_proxy_logger.debug(
|
||||
"async_data_generator: error closing response stream: %s",
|
||||
e,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
|
||||
request=request,
|
||||
request_data=request_data, # any-ok: untyped stream
|
||||
response=response, # any-ok: untyped stream
|
||||
stream_completed=stream_completed,
|
||||
client_disconnected=client_disconnected,
|
||||
)
|
||||
|
||||
|
||||
def select_data_generator(
|
||||
response, user_api_key_dict: UserAPIKeyAuth, request_data: dict
|
||||
response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_data: dict,
|
||||
request: Request | None = None,
|
||||
):
|
||||
return async_data_generator(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7250,15 +7324,53 @@ class ProxyStartupEvent:
|
|||
if _use_redis_transaction_buffer and redis_usage_cache is None:
|
||||
raise ValueError(
|
||||
"`use_redis_transaction_buffer` is enabled in general_settings "
|
||||
"but no Redis cache is configured. This will cause spend updates "
|
||||
"but no Redis is configured. This will cause spend updates "
|
||||
"to not be tracked. Add a Redis cache in litellm_settings:\n\n"
|
||||
"litellm_settings:\n"
|
||||
" cache: true\n"
|
||||
" cache_params:\n"
|
||||
" type: redis\n"
|
||||
" url: os.environ/REDIS_URL\n"
|
||||
" url: os.environ/REDIS_URL\n\n"
|
||||
"or set REDIS_* environment variables (e.g. REDIS_HOST, "
|
||||
"REDIS_PORT, REDIS_PASSWORD, or REDIS_URL) to use a standalone "
|
||||
"Redis for the transaction buffer."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_transaction_buffer_redis_cache(
|
||||
general_settings: dict,
|
||||
) -> RedisCache | None:
|
||||
"""
|
||||
Builds a standalone Redis cache from REDIS_* environment variables so
|
||||
use_redis_transaction_buffer can run when the proxy cache backend is not
|
||||
Redis (e.g. disk, s3).
|
||||
|
||||
Returns None when the buffer is disabled, or when no Redis host or url
|
||||
is set in the environment.
|
||||
"""
|
||||
from litellm._redis import _redis_kwargs_from_environment
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
_use_redis_transaction_buffer: bool | str | None = (
|
||||
general_settings.get( # any-ok: untyped stream
|
||||
"use_redis_transaction_buffer", False
|
||||
)
|
||||
)
|
||||
if isinstance(_use_redis_transaction_buffer, str):
|
||||
_use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer)
|
||||
|
||||
if not _use_redis_transaction_buffer:
|
||||
return None
|
||||
|
||||
redis_env_kwargs = _redis_kwargs_from_environment() # any-ok: untyped stream
|
||||
if (
|
||||
"host" not in redis_env_kwargs # any-ok: untyped stream
|
||||
and "url" not in redis_env_kwargs # any-ok: untyped stream
|
||||
):
|
||||
return None
|
||||
|
||||
return RedisCache(**redis_env_kwargs) # any-ok: untyped stream
|
||||
|
||||
@classmethod
|
||||
async def _initialize_semantic_tool_filter(
|
||||
cls,
|
||||
|
|
@ -8609,6 +8721,7 @@ async def chat_completion(
|
|||
response=_streaming_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
|
|
@ -8643,6 +8756,7 @@ async def chat_completion(
|
|||
response=_streaming_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
|
|
@ -8791,6 +8905,7 @@ async def completion(
|
|||
response=_streaming_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=_data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
|
|
@ -8837,6 +8952,7 @@ async def completion(
|
|||
response=_streaming_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=data,
|
||||
request=request,
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
|
|
@ -13309,6 +13425,7 @@ async def async_queue_request(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1243,6 +1243,16 @@
|
|||
"provider_display_name": "Google AI Studio",
|
||||
"litellm_provider": "gemini",
|
||||
"credential_fields": [
|
||||
{
|
||||
"key": "api_base",
|
||||
"label": "API Base",
|
||||
"placeholder": "https://generativelanguage.googleapis.com/v1beta",
|
||||
"tooltip": "Leave blank to let LiteLLM pick the right Gemini API version automatically (v1alpha for Gemini 3+ models, v1beta otherwise). Override only when fronting Gemini through a custom gateway; if you do, include the version prefix (e.g. /v1beta) but not the trailing slash. LiteLLM appends '/models/{model}:generateContent'.",
|
||||
"required": false,
|
||||
"field_type": "text",
|
||||
"options": null,
|
||||
"default_value": null
|
||||
},
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "API Key",
|
||||
|
|
|
|||
|
|
@ -244,9 +244,17 @@ def _check_non_standard_fallback_format(fallbacks: Optional[List[Any]]) -> bool:
|
|||
if all(isinstance(item, str) for item in fallbacks):
|
||||
return True
|
||||
elif all(isinstance(item, dict) for item in fallbacks):
|
||||
for key in LiteLLMParamsTypedDict.__annotations__.keys():
|
||||
if key in fallbacks[0].keys():
|
||||
return True
|
||||
for item in fallbacks: # any-ok: untyped config
|
||||
for (
|
||||
key
|
||||
) in (
|
||||
LiteLLMParamsTypedDict.__annotations__.keys() # any-ok: untyped config
|
||||
):
|
||||
if key in item: # any-ok: untyped config
|
||||
# If the value is a list, it's likely a standard fallback model group mapping
|
||||
# (e.g. {"model": ["backup"]}) rather than a parameter override.
|
||||
if not isinstance(item[key], list): # any-ok: untyped config
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.proxy._types import KeyManagementSystem
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings
|
||||
|
||||
from .base_secret_manager import BaseSecretManager
|
||||
|
||||
|
|
@ -43,6 +44,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
aws_profile_name: Optional[str] = None,
|
||||
aws_web_identity_token: Optional[str] = None,
|
||||
aws_sts_endpoint: Optional[str] = None,
|
||||
replica_regions: list[str] | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
BaseSecretManager.__init__(self, **kwargs)
|
||||
|
|
@ -56,6 +58,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
self.aws_profile_name = aws_profile_name
|
||||
self.aws_web_identity_token = aws_web_identity_token
|
||||
self.aws_sts_endpoint = aws_sts_endpoint
|
||||
self.replica_regions: list[str] = replica_regions or []
|
||||
|
||||
@classmethod
|
||||
def validate_environment(cls):
|
||||
|
|
@ -75,7 +78,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
def load_aws_secret_manager(
|
||||
cls,
|
||||
use_aws_secret_manager: Optional[bool],
|
||||
key_management_settings: Optional[Any] = None,
|
||||
key_management_settings: KeyManagementSettings | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize AWSSecretsManagerV2 with settings from key_management_settings
|
||||
|
|
@ -110,6 +113,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
"aws_sts_endpoint": getattr(
|
||||
key_management_settings, "aws_sts_endpoint", None
|
||||
),
|
||||
"replica_regions": key_management_settings.replica_regions,
|
||||
}
|
||||
# Remove None values
|
||||
aws_kwargs = {k: v for k, v in aws_kwargs.items() if v is not None}
|
||||
|
|
@ -316,6 +320,90 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
params={"timeout": timeout},
|
||||
)
|
||||
|
||||
try:
|
||||
response = await async_client.post( # any-ok: untyped httpx
|
||||
url=endpoint_url,
|
||||
headers=headers, # any-ok: untyped httpx
|
||||
data=body.decode("utf-8"), # any-ok: untyped httpx
|
||||
)
|
||||
response.raise_for_status() # any-ok: untyped httpx
|
||||
create_response = response.json() # any-ok: untyped httpx
|
||||
except httpx.HTTPStatusError as err:
|
||||
raise ValueError(f"HTTP error occurred: {err.response.text}")
|
||||
except httpx.TimeoutException:
|
||||
raise ValueError("Timeout error occurred")
|
||||
|
||||
if self.replica_regions:
|
||||
try:
|
||||
await self.async_replicate_secret(
|
||||
secret_name=secret_name,
|
||||
replica_regions=self.replica_regions,
|
||||
optional_params=optional_params, # any-ok: untyped httpx
|
||||
timeout=timeout,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Replicated secret '%s' to regions: %s",
|
||||
secret_name,
|
||||
self.replica_regions,
|
||||
)
|
||||
except Exception as replication_err: # noqa: BLE001
|
||||
verbose_logger.warning(
|
||||
"Failed to replicate secret '%s' to regions %s: %s — key was created successfully.",
|
||||
secret_name,
|
||||
self.replica_regions,
|
||||
str(replication_err),
|
||||
)
|
||||
|
||||
return create_response # any-ok: untyped httpx
|
||||
|
||||
async def async_replicate_secret(
|
||||
self,
|
||||
secret_name: str,
|
||||
replica_regions: list[str],
|
||||
optional_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Replicate a secret to additional AWS regions using ReplicateSecretToRegions.
|
||||
|
||||
Called after a successful CreateSecret when replica_regions is configured.
|
||||
Replication is best-effort — callers should not depend on this for correctness.
|
||||
|
||||
Args:
|
||||
secret_name: Name or ARN of the secret to replicate
|
||||
replica_regions: List of target AWS region names, e.g. ["us-west-2"]
|
||||
optional_params: Additional AWS parameters
|
||||
timeout: Request timeout
|
||||
|
||||
Returns:
|
||||
dict: AWS response, or {} if replica_regions is empty
|
||||
"""
|
||||
if not replica_regions:
|
||||
return {}
|
||||
|
||||
verbose_logger.info(
|
||||
"ReplicateSecretToRegions called for secret '%s' in regions %s",
|
||||
secret_name,
|
||||
replica_regions,
|
||||
)
|
||||
|
||||
data: dict[str, object] = {
|
||||
"SecretId": secret_name,
|
||||
"AddReplicaRegions": [{"Region": r} for r in replica_regions],
|
||||
}
|
||||
|
||||
endpoint_url, headers, body = self._prepare_request( # any-ok: untyped httpx
|
||||
action="ReplicateSecretToRegions",
|
||||
secret_name=secret_name,
|
||||
optional_params=optional_params,
|
||||
request_data=data,
|
||||
)
|
||||
|
||||
async_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.SecretManager,
|
||||
params={"timeout": timeout}, # any-ok: untyped httpx
|
||||
)
|
||||
|
||||
try:
|
||||
response = await async_client.post(
|
||||
url=endpoint_url, headers=headers, data=body.decode("utf-8")
|
||||
|
|
|
|||
|
|
@ -11,5 +11,6 @@ class UiDiscoveryEndpoints(BaseModel):
|
|||
auto_redirect_to_sso: bool
|
||||
admin_ui_disabled: bool
|
||||
sso_configured: bool
|
||||
hide_default_credentials_hint: bool = False
|
||||
is_control_plane: bool = False
|
||||
workers: List[WorkerRegistryEntry] = []
|
||||
|
|
|
|||
|
|
@ -72,3 +72,12 @@ class KeyManagementSettings(LiteLLMPydanticObjectBase):
|
|||
|
||||
aws_sts_endpoint: Optional[str] = None
|
||||
"""Custom STS endpoint URL (useful for VPC endpoints or testing)"""
|
||||
|
||||
replica_regions: Optional[List[str]] = None
|
||||
"""
|
||||
Optional list of additional AWS regions to replicate secrets to after CreateSecret.
|
||||
Uses the AWS Secrets Manager ReplicateSecretToRegions API. Replication is
|
||||
best-effort — failure to replicate does not fail key creation.
|
||||
Example: ["us-west-2", "eu-west-1"]
|
||||
Only applies when key_management_system is "aws_secret_manager".
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -196,7 +196,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
float
|
||||
] # OpenAI priority service tier pricing
|
||||
cache_read_input_token_cost_above_200k_tokens: Optional[float]
|
||||
cache_read_input_token_cost_above_200k_tokens_priority: Optional[float]
|
||||
cache_read_input_token_cost_above_272k_tokens: Optional[float]
|
||||
cache_read_input_token_cost_above_272k_tokens_priority: Optional[float]
|
||||
cache_read_input_token_cost_above_512k_tokens: Optional[float]
|
||||
input_cost_per_character: Optional[float] # only for vertex ai models
|
||||
input_cost_per_audio_token: Optional[float]
|
||||
|
|
@ -204,9 +206,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_token_above_200k_tokens: Optional[
|
||||
float
|
||||
] # only for vertex ai gemini-2.5-pro models
|
||||
input_cost_per_token_above_200k_tokens_priority: Optional[float]
|
||||
input_cost_per_token_above_272k_tokens: Optional[
|
||||
float
|
||||
] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input
|
||||
input_cost_per_token_above_272k_tokens_priority: Optional[float]
|
||||
input_cost_per_token_above_512k_tokens: Optional[
|
||||
float
|
||||
] # MiniMax-M3: prompts >512K priced at 2x input
|
||||
|
|
@ -240,9 +244,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
output_cost_per_token_above_200k_tokens: Optional[
|
||||
float
|
||||
] # only for vertex ai gemini-2.5-pro models
|
||||
output_cost_per_token_above_200k_tokens_priority: Optional[float]
|
||||
output_cost_per_token_above_272k_tokens: Optional[
|
||||
float
|
||||
] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output
|
||||
output_cost_per_token_above_272k_tokens_priority: Optional[float]
|
||||
output_cost_per_token_above_512k_tokens: Optional[
|
||||
float
|
||||
] # MiniMax-M3: prompts >512K priced at 2x output
|
||||
|
|
@ -3093,6 +3099,8 @@ class CustomPricingLiteLLMParams(BaseModel):
|
|||
cache_read_input_token_cost_flex: Optional[float] = None
|
||||
cache_read_input_token_cost_priority: Optional[float] = None
|
||||
cache_read_input_token_cost_above_200k_tokens: Optional[float] = None
|
||||
cache_read_input_token_cost_above_200k_tokens_priority: Optional[float] = None
|
||||
cache_read_input_token_cost_above_272k_tokens_priority: Optional[float] = None
|
||||
cache_read_input_audio_token_cost: Optional[float] = None
|
||||
input_cost_per_character: Optional[float] = None
|
||||
input_cost_per_character_above_128k_tokens: Optional[float] = None
|
||||
|
|
@ -3100,6 +3108,8 @@ class CustomPricingLiteLLMParams(BaseModel):
|
|||
input_cost_per_token_cache_hit: Optional[float] = None
|
||||
input_cost_per_token_above_128k_tokens: Optional[float] = None
|
||||
input_cost_per_token_above_200k_tokens: Optional[float] = None
|
||||
input_cost_per_token_above_200k_tokens_priority: Optional[float] = None
|
||||
input_cost_per_token_above_272k_tokens_priority: Optional[float] = None
|
||||
input_cost_per_query: Optional[float] = None
|
||||
input_cost_per_image: Optional[float] = None
|
||||
input_cost_per_image_above_128k_tokens: Optional[float] = None
|
||||
|
|
@ -3117,6 +3127,8 @@ class CustomPricingLiteLLMParams(BaseModel):
|
|||
output_cost_per_audio_token: Optional[float] = None
|
||||
output_cost_per_token_above_128k_tokens: Optional[float] = None
|
||||
output_cost_per_token_above_200k_tokens: Optional[float] = None
|
||||
output_cost_per_token_above_200k_tokens_priority: Optional[float] = None
|
||||
output_cost_per_token_above_272k_tokens_priority: Optional[float] = None
|
||||
output_cost_per_character_above_128k_tokens: Optional[float] = None
|
||||
output_cost_per_image: Optional[float] = None
|
||||
output_cost_per_image_token: Optional[float] = None
|
||||
|
|
|
|||
|
|
@ -6043,9 +6043,15 @@ def _get_model_info_helper(
|
|||
cache_read_input_token_cost_above_200k_tokens=_model_info.get(
|
||||
"cache_read_input_token_cost_above_200k_tokens", None
|
||||
),
|
||||
cache_read_input_token_cost_above_200k_tokens_priority=_model_info.get( # any-ok: untyped cost map
|
||||
"cache_read_input_token_cost_above_200k_tokens_priority", None
|
||||
),
|
||||
cache_read_input_token_cost_above_272k_tokens=_model_info.get(
|
||||
"cache_read_input_token_cost_above_272k_tokens", None
|
||||
),
|
||||
cache_read_input_token_cost_above_272k_tokens_priority=_model_info.get( # any-ok: untyped cost map
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority", None
|
||||
),
|
||||
cache_read_input_token_cost_above_512k_tokens=_model_info.get(
|
||||
"cache_read_input_token_cost_above_512k_tokens", None
|
||||
),
|
||||
|
|
@ -6067,9 +6073,15 @@ def _get_model_info_helper(
|
|||
input_cost_per_token_above_200k_tokens=_model_info.get(
|
||||
"input_cost_per_token_above_200k_tokens", None
|
||||
),
|
||||
input_cost_per_token_above_200k_tokens_priority=_model_info.get( # any-ok: untyped cost map
|
||||
"input_cost_per_token_above_200k_tokens_priority", None
|
||||
),
|
||||
input_cost_per_token_above_272k_tokens=_model_info.get(
|
||||
"input_cost_per_token_above_272k_tokens", None
|
||||
),
|
||||
input_cost_per_token_above_272k_tokens_priority=_model_info.get( # any-ok: untyped cost map
|
||||
"input_cost_per_token_above_272k_tokens_priority", None
|
||||
),
|
||||
input_cost_per_token_above_512k_tokens=_model_info.get(
|
||||
"input_cost_per_token_above_512k_tokens", None
|
||||
),
|
||||
|
|
@ -6125,9 +6137,15 @@ def _get_model_info_helper(
|
|||
output_cost_per_token_above_200k_tokens=_model_info.get(
|
||||
"output_cost_per_token_above_200k_tokens", None
|
||||
),
|
||||
output_cost_per_token_above_200k_tokens_priority=_model_info.get( # any-ok: untyped cost map
|
||||
"output_cost_per_token_above_200k_tokens_priority", None
|
||||
),
|
||||
output_cost_per_token_above_272k_tokens=_model_info.get(
|
||||
"output_cost_per_token_above_272k_tokens", None
|
||||
),
|
||||
output_cost_per_token_above_272k_tokens_priority=_model_info.get( # any-ok: untyped cost map
|
||||
"output_cost_per_token_above_272k_tokens_priority", None
|
||||
),
|
||||
output_cost_per_token_above_512k_tokens=_model_info.get(
|
||||
"output_cost_per_token_above_512k_tokens", None
|
||||
),
|
||||
|
|
|
|||
|
|
@ -2528,6 +2528,100 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"azure_ai/gpt-5.5": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"source": "https://ai.azure.com/catalog/models/gpt-5.5",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": 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_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false
|
||||
},
|
||||
"azure_ai/gpt-5.5-2026-04-23": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 1e-06,
|
||||
"cache_read_input_token_cost_priority": 1e-06,
|
||||
"cache_read_input_token_cost_above_272k_tokens_priority": 2e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_272k_tokens": 1e-05,
|
||||
"input_cost_per_token_priority": 1e-05,
|
||||
"input_cost_per_token_above_272k_tokens_priority": 2e-05,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 1050000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"output_cost_per_token_above_272k_tokens": 4.5e-05,
|
||||
"output_cost_per_token_priority": 6e-05,
|
||||
"output_cost_per_token_above_272k_tokens_priority": 9e-05,
|
||||
"source": "https://ai.azure.com/catalog/models/gpt-5.5",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": 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_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_web_search": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": false
|
||||
},
|
||||
"azure_ai/gpt-5.4": {
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 5e-07,
|
||||
|
|
@ -10068,6 +10162,8 @@
|
|||
},
|
||||
"claude-sonnet-4-5": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
|
|
@ -10097,6 +10193,8 @@
|
|||
},
|
||||
"claude-sonnet-4-5-20250929": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
|
|
@ -10127,6 +10225,7 @@
|
|||
},
|
||||
"claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "anthropic",
|
||||
|
|
@ -10155,6 +10254,8 @@
|
|||
},
|
||||
"claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
"cache_creation_input_token_cost_above_1hr": 6e-06,
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05,
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
|
|
@ -25103,6 +25204,21 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"mistral/mistral-medium-3-5": {
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "mistral",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7.5e-06,
|
||||
"source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04",
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"mistral/mistral-small": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "mistral",
|
||||
|
|
@ -42830,4 +42946,105 @@
|
|||
"supports_reasoning": true,
|
||||
"source": "https://serverless.tensormesh.ai/v1/models/openrouter"
|
||||
}
|
||||
}
|
||||
,
|
||||
"deepseek-v4-flash": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"input_cost_per_token_cache_hit": 2.8e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"deepseek-v4-pro": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 3.625e-09,
|
||||
"input_cost_per_token": 4.35e-07,
|
||||
"input_cost_per_token_cache_hit": 3.625e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.7e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"deepseek/deepseek-v4-flash": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"input_cost_per_token_cache_hit": 2.8e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
},
|
||||
"deepseek/deepseek-v4-pro": {
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"cache_read_input_token_cost": 3.625e-09,
|
||||
"input_cost_per_token": 4.35e-07,
|
||||
"input_cost_per_token_cache_hit": 3.625e-09,
|
||||
"litellm_provider": "deepseek",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.7e-07,
|
||||
"source": "https://api-docs.deepseek.com/quick_start/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ try:
|
|||
from mypy.find_sources import create_source_list
|
||||
from mypy.fscache import FileSystemCache
|
||||
from mypy.modulefinder import BuildSource
|
||||
from mypy.nodes import AssignmentStmt, Expression, NameExpr, Node
|
||||
from mypy.nodes import AssignmentStmt, Expression, NameExpr, Node, TempNode
|
||||
from mypy.options import Options
|
||||
from mypy.types import (
|
||||
AnyType,
|
||||
|
|
@ -134,6 +134,19 @@ _HARMLESS_ANY = frozenset(
|
|||
# against ExtendedTraverserVisitor across the full grammar (see commit notes).
|
||||
_NON_SYNTACTIC_ATTRS = frozenset({"node", "info"})
|
||||
|
||||
# Awaitable / coroutine / generator instances carry synthetic `Any` in their
|
||||
# send (and, for coroutines, yield) protocol slots: `async def f() -> float`
|
||||
# produces `Coroutine[Any, Any, float]`, so the bare call expression `f()` would
|
||||
# be flagged even though the awaited value is a clean `float`. Only the args that
|
||||
# hold a value the caller observes (the awaited result, the yielded item) are
|
||||
# meaningful; a real `Any` there -- e.g. a coroutine that returns `Any` -- is
|
||||
# still caught because that index is still checked.
|
||||
_SYNTHETIC_SEND_YIELD_VALUE_ARGS: dict[str, tuple[int, ...]] = {
|
||||
"typing.Coroutine": (2,),
|
||||
"typing.Generator": (0, 2),
|
||||
"typing.AsyncGenerator": (0,),
|
||||
}
|
||||
|
||||
|
||||
class Violation(NamedTuple):
|
||||
path: Path
|
||||
|
|
@ -168,6 +181,12 @@ def contains_any(t: Type, _seen: set[int] | None = None) -> bool:
|
|||
if isinstance(p, UnionType):
|
||||
return any(contains_any(item, seen) for item in p.items)
|
||||
if isinstance(p, Instance):
|
||||
value_arg_indices = _SYNTHETIC_SEND_YIELD_VALUE_ARGS.get(p.type.fullname)
|
||||
if value_arg_indices is not None:
|
||||
return any(
|
||||
index < len(p.args) and contains_any(p.args[index], seen)
|
||||
for index in value_arg_indices
|
||||
)
|
||||
return any(contains_any(arg, seen) for arg in p.args)
|
||||
if isinstance(p, TupleType):
|
||||
return any(contains_any(item, seen) for item in p.items)
|
||||
|
|
@ -224,7 +243,11 @@ def find_any_in_tree(tree: Node, idmap: dict[int, Type]) -> list[tuple[int, int,
|
|||
exprs, skip_lvalues = _walk_file(tree)
|
||||
findings: list[tuple[int, int, str]] = []
|
||||
for expr in exprs:
|
||||
if id(expr) in skip_lvalues:
|
||||
# A TempNode is mypy's synthetic placeholder for a position with no real
|
||||
# expression -- e.g. the rvalue of an annotation-only `field: T` in a
|
||||
# TypedDict / class body, whose `special_form` `Any` is not a value the
|
||||
# author wrote. It never corresponds to a runtime value, so skip it.
|
||||
if id(expr) in skip_lvalues or isinstance(expr, TempNode):
|
||||
continue
|
||||
t = idmap.get(id(expr))
|
||||
if t is not None and contains_any(t):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,49 @@
|
|||
"""
|
||||
Test that check_and_fix_namespace handles None key gracefully.
|
||||
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/30424
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
|
||||
def test_check_and_fix_namespace_with_none_key():
|
||||
"""When key is None, check_and_fix_namespace should return None without raising."""
|
||||
cache = MagicMock(spec=RedisCache)
|
||||
cache.namespace = "litellm"
|
||||
# Call the real method
|
||||
result = RedisCache.check_and_fix_namespace(cache, key=None)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_check_and_fix_namespace_with_none_key_no_namespace():
|
||||
"""When key is None and namespace is None, should return None without raising."""
|
||||
cache = MagicMock(spec=RedisCache)
|
||||
cache.namespace = None
|
||||
result = RedisCache.check_and_fix_namespace(cache, key=None)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_check_and_fix_namespace_with_valid_key():
|
||||
"""Normal behavior: prefix key with namespace if not already prefixed."""
|
||||
cache = MagicMock(spec=RedisCache)
|
||||
cache.namespace = "litellm"
|
||||
result = RedisCache.check_and_fix_namespace(cache, key="my_key")
|
||||
assert result == "litellm:my_key"
|
||||
|
||||
|
||||
def test_check_and_fix_namespace_with_already_prefixed_key():
|
||||
"""If key already starts with namespace, don't double-prefix."""
|
||||
cache = MagicMock(spec=RedisCache)
|
||||
cache.namespace = "litellm"
|
||||
result = RedisCache.check_and_fix_namespace(cache, key="litellm:my_key")
|
||||
assert result == "litellm:my_key"
|
||||
|
||||
|
||||
def test_check_and_fix_namespace_no_namespace():
|
||||
"""When namespace is None, return key as-is."""
|
||||
cache = MagicMock(spec=RedisCache)
|
||||
cache.namespace = None
|
||||
result = RedisCache.check_and_fix_namespace(cache, key="my_key")
|
||||
assert result == "my_key"
|
||||
|
|
@ -1573,3 +1573,72 @@ def test_data_residency_composes_with_service_tier(_local_model_cost_map):
|
|||
|
||||
assert priority_base_total > 0
|
||||
assert priority_eu_total == pytest.approx(priority_base_total * 1.10, rel=1e-9)
|
||||
|
||||
|
||||
def test_priority_service_tier_above_threshold_uses_priority_tier_rates_for_cached_tokens(
|
||||
_local_model_cost_map,
|
||||
):
|
||||
"""Regression: for a model that publishes both service_tier and above_threshold rate
|
||||
variants, a priority request over the threshold must bill cached tokens at
|
||||
cache_read_input_token_cost_above_200k_tokens_priority (and analogously for
|
||||
input/output above-threshold), not the standard above-threshold rate."""
|
||||
usage = Usage(
|
||||
prompt_tokens=250_000,
|
||||
completion_tokens=1_000,
|
||||
total_tokens=251_000,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=200_000, text_tokens=50_000
|
||||
),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=1_000),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="gemini-3-pro-preview",
|
||||
usage=usage,
|
||||
custom_llm_provider="gemini",
|
||||
service_tier="priority",
|
||||
)
|
||||
|
||||
# gemini-3-pro-preview priority + above_200k rates from the pricing JSON:
|
||||
# input 7.2e-6, output 3.24e-5, cache_read 7.2e-7
|
||||
expected_prompt = 50_000 * 7.2e-6 + 200_000 * 7.2e-7
|
||||
expected_completion = 1_000 * 3.24e-5
|
||||
assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
|
||||
assert completion_cost == pytest.approx(expected_completion, rel=1e-9)
|
||||
|
||||
|
||||
def test_priority_service_tier_above_threshold_falls_back_to_standard_for_cache_creation(
|
||||
_local_model_cost_map,
|
||||
):
|
||||
"""Regression: priority requests against models that publish standard above-threshold
|
||||
cache_creation rates but no priority variant must fall back to the standard
|
||||
above-threshold rate, not the priority-base rate. vertex_ai/claude-sonnet-4-5
|
||||
has cache_creation_input_token_cost_above_200k_tokens but no _priority sibling."""
|
||||
usage = Usage(
|
||||
prompt_tokens=350_000,
|
||||
completion_tokens=1_000,
|
||||
total_tokens=351_000,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=200_000,
|
||||
cache_creation_tokens=100_000,
|
||||
text_tokens=50_000,
|
||||
),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=1_000),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="vertex_ai/claude-sonnet-4-5",
|
||||
usage=usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
service_tier="priority",
|
||||
)
|
||||
|
||||
# vertex_ai/claude-sonnet-4-5 above_200k (no _priority variants):
|
||||
# input 6e-6, output 2.25e-5, cache_read 6e-7, cache_creation 7.5e-6
|
||||
# text 50_000 * 6e-6 = 0.30
|
||||
# cache_read 200_000 * 6e-7 = 0.12
|
||||
# cache_creation 100_000 * 7.5e-6 = 0.75
|
||||
expected_prompt = 50_000 * 6e-6 + 200_000 * 6e-7 + 100_000 * 7.5e-6
|
||||
expected_completion = 1_000 * 2.25e-5
|
||||
assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
|
||||
assert completion_cost == pytest.approx(expected_completion, rel=1e-9)
|
||||
|
|
|
|||
|
|
@ -3115,6 +3115,71 @@ class TestFirstApiCallStartTimeSetOnce:
|
|||
assert user_meta == {}
|
||||
|
||||
|
||||
def test_get_error_information_for_logging_payload_ignores_spoofed_disconnect_without_flag():
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
baseline = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=ValueError("provider failure"),
|
||||
)
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={
|
||||
"error_information": {
|
||||
"error_code": "499",
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
}
|
||||
},
|
||||
original_exception=ValueError("provider failure"),
|
||||
error_str="provider failure",
|
||||
)
|
||||
)
|
||||
assert error_information == baseline
|
||||
assert error_str == "provider failure"
|
||||
|
||||
|
||||
def test_get_error_information_for_logging_payload_client_disconnect():
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
custom_error = {
|
||||
"error_code": "499",
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
}
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={"client_disconnected": True, "error_information": custom_error},
|
||||
original_exception=None,
|
||||
error_str=None,
|
||||
)
|
||||
)
|
||||
assert error_information == custom_error
|
||||
assert error_str == "Client disconnected the request"
|
||||
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={"client_disconnected": True},
|
||||
original_exception=None,
|
||||
error_str="existing error",
|
||||
)
|
||||
)
|
||||
assert error_information["error_code"] == "499"
|
||||
assert error_str == "existing error"
|
||||
|
||||
baseline = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=None,
|
||||
)
|
||||
error_information, error_str = (
|
||||
StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
metadata={},
|
||||
original_exception=None,
|
||||
error_str=None,
|
||||
)
|
||||
)
|
||||
assert error_information == baseline
|
||||
assert error_str is None
|
||||
|
||||
|
||||
def test_get_error_information_proxy_exception_preserves_message():
|
||||
"""ProxyException keeps its text in ``.message`` (str() was empty pre-fix),
|
||||
so error_information must still surface the message and code."""
|
||||
|
|
|
|||
|
|
@ -523,7 +523,6 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding
|
||||
|
||||
|
||||
# Clear the cache at module load to ensure clean state
|
||||
_select_tokenizer_helper.cache_clear()
|
||||
|
||||
|
|
@ -1010,3 +1009,64 @@ def test_token_counter_with_thinking_content():
|
|||
assert (
|
||||
tokens_no_thinking < 15
|
||||
), f"Expected minimal token count for empty thinking block, got {tokens_no_thinking}"
|
||||
|
||||
|
||||
def test_token_counter_with_tool_reference_block():
|
||||
"""
|
||||
Regression test: a message containing an Anthropic tool-search
|
||||
`tool_reference` content block must NOT raise.
|
||||
|
||||
Before the fix, token_counter raised
|
||||
`Invalid content item type: tool_reference`. On the streaming
|
||||
anthropic_messages proxy path this nulled response_cost and caused the
|
||||
SpendLogs row to be dropped, silently undercounting cost. token_counter
|
||||
must instead count the referenced tool name and return a positive count.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Let me look up the right tool."},
|
||||
{"type": "tool_reference", "tool_name": "search_knowledge_base"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# Must not raise, and must produce a positive token count.
|
||||
tokens = token_counter_new(
|
||||
model="anthropic/claude-sonnet-4-5-20250929", messages=messages
|
||||
)
|
||||
assert tokens > 0, f"Expected positive token count, got {tokens}"
|
||||
|
||||
# A tool_reference with no/empty tool_name must also be handled gracefully.
|
||||
messages_empty = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_reference", "tool_name": ""}],
|
||||
}
|
||||
]
|
||||
tokens_empty = token_counter_new(
|
||||
model="anthropic/claude-sonnet-4-5-20250929", messages=messages_empty
|
||||
)
|
||||
assert tokens_empty >= 0
|
||||
|
||||
|
||||
def test_count_content_list_rejects_unknown_type():
|
||||
"""
|
||||
An unrecognized content block type must raise, and the error message must
|
||||
enumerate the supported types (including `tool_reference`). This pins the
|
||||
catch-all contract so a future block type isn't silently dropped.
|
||||
"""
|
||||
from litellm.litellm_core_utils.token_counter import _count_content_list
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
_count_content_list(
|
||||
count_function=len,
|
||||
content_list=[{"type": "totally_unknown_block"}],
|
||||
use_default_image_token_count=False,
|
||||
default_token_count=None,
|
||||
)
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "Invalid content item type: totally_unknown_block" in message
|
||||
assert "tool_reference" in message
|
||||
|
|
|
|||
|
|
@ -0,0 +1,131 @@
|
|||
"""
|
||||
Integration / regression tests for Anthropic tool-search (`tool_reference`)
|
||||
content blocks on the cost-calculation and streaming-assembly paths used by
|
||||
Claude Code.
|
||||
|
||||
Claude Code's tool-search feature emits assistant content blocks of the form
|
||||
``{"type": "tool_reference", "tool_name": ...}`` -- a lightweight pointer to a
|
||||
deferred tool. Before the fix, `token_counter` did not recognise this block
|
||||
type and raised ``Invalid content item type: tool_reference``.
|
||||
|
||||
Why this matters (the bug these tests guard against):
|
||||
|
||||
* On the cost path, that exception propagates out of ``completion_cost`` ->
|
||||
``response_cost_calculator``. The proxy logging layer catches it and nulls
|
||||
``response_cost``; the spend-tracking callback then skips the request, so
|
||||
the entire SpendLogs row is dropped. The request succeeds for the caller
|
||||
but the spend is silently never recorded -- a cost undercount on ALL
|
||||
tool-search traffic.
|
||||
|
||||
* On the streaming-assembly path, ``stream_chunk_builder`` recomputes the
|
||||
prompt tokens from the request messages when the provider stream does not
|
||||
carry usage. The same exception there was swallowed and prompt tokens
|
||||
silently collapsed to 0 -- a quieter undercount of the same traffic.
|
||||
|
||||
These tests exercise the real public entry points (not the private
|
||||
``_count_content_list`` helper) so the whole chain is covered end to end.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
import litellm
|
||||
from litellm import stream_chunk_builder
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
ANTHROPIC_MODEL = "anthropic/claude-sonnet-4-5-20250929"
|
||||
|
||||
# Mirrors a Claude Code tool-search turn: a normal text block followed by a
|
||||
# `tool_reference` pointer to a deferred tool.
|
||||
TOOL_SEARCH_MESSAGES = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Let me look up the right tool."},
|
||||
{"type": "tool_reference", "tool_name": "search_knowledge_base"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_completion_cost_with_tool_reference_records_spend():
|
||||
"""
|
||||
``completion_cost`` must return a real, positive cost for messages that
|
||||
contain a tool-search ``tool_reference`` block.
|
||||
|
||||
This is the exact chain that fails on the streaming anthropic_messages
|
||||
proxy path: before the fix ``completion_cost`` raised, the logging layer
|
||||
caught the exception and set ``response_cost = None``, and the spend
|
||||
callback then dropped the SpendLogs row. A positive cost here means the
|
||||
row is recorded instead of silently dropped.
|
||||
"""
|
||||
cost = litellm.completion_cost(model=ANTHROPIC_MODEL, messages=TOOL_SEARCH_MESSAGES)
|
||||
|
||||
assert cost is not None, "response_cost is None -> SpendLogs row would be dropped"
|
||||
assert cost > 0, f"Expected a positive cost for tool-search traffic, got {cost}"
|
||||
|
||||
|
||||
def test_completion_cost_with_empty_tool_name_records_spend():
|
||||
"""A ``tool_reference`` with an empty/missing ``tool_name`` must also cost
|
||||
out cleanly rather than raising and nulling the spend."""
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_reference", "tool_name": ""}],
|
||||
}
|
||||
]
|
||||
|
||||
cost = litellm.completion_cost(model=ANTHROPIC_MODEL, messages=messages)
|
||||
|
||||
assert cost is not None
|
||||
assert cost >= 0
|
||||
|
||||
|
||||
def test_stream_chunk_builder_counts_prompt_tokens_for_tool_reference():
|
||||
"""
|
||||
On the streaming-assembly path used by Claude Code, when the provider
|
||||
stream carries no prompt-token usage, ``stream_chunk_builder`` recomputes
|
||||
prompt tokens from the request messages via ``token_counter``.
|
||||
|
||||
With a ``tool_reference`` block in those messages the count must be
|
||||
positive. Before the fix the underlying ``token_counter`` call raised and
|
||||
the assembler swallowed it, collapsing ``prompt_tokens`` to 0 -- a silent
|
||||
undercount of every tool-search request.
|
||||
"""
|
||||
model = "claude-sonnet-4-5-20250929"
|
||||
# Chunks deliberately carry no usage, forcing the prompt-token fallback.
|
||||
chunks = [
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-tool-search",
|
||||
created=1700000000,
|
||||
model=model,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(content="Searching...", role="assistant"),
|
||||
)
|
||||
],
|
||||
),
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-tool-search",
|
||||
created=1700000000,
|
||||
model=model,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason="stop", index=0, delta=Delta(content="")
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
response = stream_chunk_builder(chunks, messages=TOOL_SEARCH_MESSAGES)
|
||||
|
||||
assert response is not None
|
||||
assert (
|
||||
response.usage.prompt_tokens > 0
|
||||
), "prompt_tokens collapsed to 0 -> tool-search traffic silently undercounted"
|
||||
|
|
@ -169,8 +169,8 @@ def test_hosted_vllm_supports_thinking():
|
|||
|
||||
def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content():
|
||||
"""
|
||||
Test that thinking_blocks on assistant messages are converted to content
|
||||
blocks prepended before the existing content.
|
||||
Test that thinking_blocks on assistant messages are removed and content
|
||||
stays a string for vLLM compatibility.
|
||||
"""
|
||||
config = HostedVLLMChatConfig()
|
||||
messages = [
|
||||
|
|
@ -203,21 +203,15 @@ def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content():
|
|||
)
|
||||
assistant_msg = transformed["messages"][1]
|
||||
assert assistant_msg["role"] == "assistant"
|
||||
assert isinstance(assistant_msg["content"], list)
|
||||
assert assistant_msg["content"][0] == {
|
||||
"type": "thinking",
|
||||
"thinking": "Let me reason about this...",
|
||||
}
|
||||
assert assistant_msg["content"][1] == {
|
||||
"type": "text",
|
||||
"text": "Here is my answer.",
|
||||
}
|
||||
assert isinstance(assistant_msg["content"], str)
|
||||
assert assistant_msg["content"] == "Here is my answer."
|
||||
assert "thinking_blocks" not in assistant_msg
|
||||
|
||||
|
||||
def test_hosted_vllm_thinking_blocks_with_list_content():
|
||||
"""
|
||||
Test thinking_blocks prepended when assistant content is already a list.
|
||||
Test thinking_blocks are removed and assistant content list is converted
|
||||
to a string.
|
||||
"""
|
||||
config = HostedVLLMChatConfig()
|
||||
messages = [
|
||||
|
|
@ -246,19 +240,125 @@ def test_hosted_vllm_thinking_blocks_with_list_content():
|
|||
headers={},
|
||||
)
|
||||
assistant_msg = transformed["messages"][0]
|
||||
assert len(assistant_msg["content"]) == 3
|
||||
assert assistant_msg["content"][0] == {
|
||||
"type": "thinking",
|
||||
"thinking": "Step 1 reasoning",
|
||||
}
|
||||
assert assistant_msg["content"][1] == {
|
||||
"type": "thinking",
|
||||
"thinking": "Step 2 reasoning",
|
||||
}
|
||||
assert assistant_msg["content"][2] == {"type": "text", "text": "Response text"}
|
||||
assert isinstance(assistant_msg["content"], str)
|
||||
assert assistant_msg["content"] == "Response text"
|
||||
assert "thinking_blocks" not in assistant_msg
|
||||
|
||||
|
||||
def test_hosted_vllm_assistant_structured_content_is_preserved():
|
||||
config = HostedVLLMChatConfig()
|
||||
image_block = {
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "https://example.com/image.png"},
|
||||
}
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "Here is the image"}, image_block],
|
||||
},
|
||||
]
|
||||
|
||||
transformed = config.transform_request(
|
||||
model="hosted_vllm/llama-3.1-70b-instruct",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assistant_msg = transformed["messages"][0]
|
||||
assert assistant_msg["content"] == [
|
||||
{"type": "text", "text": "Here is the image"},
|
||||
image_block,
|
||||
]
|
||||
|
||||
|
||||
def test_hosted_vllm_assistant_tool_use_content_becomes_tool_calls():
|
||||
config = HostedVLLMChatConfig()
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_1",
|
||||
"name": "get_weather",
|
||||
"input": {"city": "Boston"},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
transformed = config.transform_request(
|
||||
model="hosted_vllm/llama-3.1-70b-instruct",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assistant_msg = transformed["messages"][0]
|
||||
assert assistant_msg["content"] == ""
|
||||
assert assistant_msg["tool_calls"] == [
|
||||
{
|
||||
"id": "toolu_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": json.dumps({"city": "Boston"}),
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_hosted_vllm_assistant_tool_use_does_not_duplicate_existing_tool_calls():
|
||||
config = HostedVLLMChatConfig()
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_1",
|
||||
"name": "get_weather",
|
||||
"input": {"city": "Boston"},
|
||||
}
|
||||
],
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "toolu_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": json.dumps({"city": "Boston"}),
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
transformed = config.transform_request(
|
||||
model="hosted_vllm/llama-3.1-70b-instruct",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assistant_msg = transformed["messages"][0]
|
||||
assert assistant_msg["content"] == ""
|
||||
assert assistant_msg["tool_calls"] == [
|
||||
{
|
||||
"id": "toolu_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": json.dumps({"city": "Boston"}),
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_hosted_vllm_custom_tools_are_converted_to_function_tools():
|
||||
config = HostedVLLMChatConfig()
|
||||
optional_params = config.map_openai_params(
|
||||
|
|
|
|||
|
|
@ -553,3 +553,68 @@ def test_openrouter_non_reasoning_models_do_not_add_reasoning_effort():
|
|||
)
|
||||
|
||||
assert "reasoning_effort" not in supported_params
|
||||
|
||||
|
||||
def test_openrouter_reasoning_effort_max_maps_to_xhigh():
|
||||
"""
|
||||
OpenRouter expects 'xhigh' instead of 'max' for reasoning_effort.
|
||||
"""
|
||||
config = OpenrouterConfig()
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params={"reasoning_effort": "max"},
|
||||
optional_params={},
|
||||
model="openrouter/deepseek/deepseek-r1",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["reasoning_effort"] == "xhigh"
|
||||
|
||||
|
||||
def test_openrouter_reasoning_effort_max_does_not_mutate_caller_dict():
|
||||
"""
|
||||
map_openai_params must not mutate the caller-supplied non_default_params dict.
|
||||
"""
|
||||
config = OpenrouterConfig()
|
||||
original_params = {"reasoning_effort": "max"}
|
||||
|
||||
config.map_openai_params(
|
||||
non_default_params=original_params,
|
||||
optional_params={},
|
||||
model="openrouter/deepseek/deepseek-r1",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert original_params["reasoning_effort"] == "max"
|
||||
|
||||
|
||||
def test_openrouter_reasoning_effort_xhigh_passes_through():
|
||||
"""
|
||||
reasoning_effort='xhigh' should be forwarded unchanged.
|
||||
"""
|
||||
config = OpenrouterConfig()
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params={"reasoning_effort": "xhigh"},
|
||||
optional_params={},
|
||||
model="openrouter/deepseek/deepseek-r1",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["reasoning_effort"] == "xhigh"
|
||||
|
||||
|
||||
def test_openrouter_reasoning_effort_high_passes_through():
|
||||
"""
|
||||
Non-max reasoning_effort values should be forwarded unchanged.
|
||||
"""
|
||||
config = OpenrouterConfig()
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params={"reasoning_effort": "high"},
|
||||
optional_params={},
|
||||
model="openrouter/deepseek/deepseek-r1",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["reasoning_effort"] == "high"
|
||||
|
|
|
|||
|
|
@ -4996,3 +4996,146 @@ def test_mid_stream_429_error_raises_during_iteration():
|
|||
# Verify: 429 error is properly raised
|
||||
assert exc_info.value.status_code == 429
|
||||
assert "RESOURCE_EXHAUSTED" in str(exc_info.value.message)
|
||||
|
||||
|
||||
class TestModelResponseIteratorCleanup:
|
||||
def _make_logging_obj(self):
|
||||
from unittest.mock import Mock
|
||||
|
||||
obj = Mock()
|
||||
obj.optional_params = {}
|
||||
return obj
|
||||
|
||||
def test_aclose_closes_iterator_and_response(self):
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.aclose = AsyncMock()
|
||||
|
||||
mock_iterator = MagicMock()
|
||||
mock_iterator.aclose = AsyncMock()
|
||||
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(),
|
||||
sync_stream=False,
|
||||
logging_obj=self._make_logging_obj(),
|
||||
response=mock_response,
|
||||
)
|
||||
iterator.async_response_iterator = mock_iterator
|
||||
|
||||
asyncio.run(iterator.aclose())
|
||||
|
||||
mock_iterator.aclose.assert_awaited_once()
|
||||
mock_response.aclose.assert_awaited_once()
|
||||
|
||||
def test_close_closes_iterator_and_response(self):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_iterator = MagicMock()
|
||||
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(),
|
||||
sync_stream=True,
|
||||
logging_obj=self._make_logging_obj(),
|
||||
response=mock_response,
|
||||
)
|
||||
iterator.response_iterator = mock_iterator
|
||||
|
||||
iterator.close()
|
||||
|
||||
mock_iterator.close.assert_called_once()
|
||||
mock_response.close.assert_called_once()
|
||||
|
||||
def test_aclose_without_response_does_not_raise(self):
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
mock_iterator = MagicMock()
|
||||
mock_iterator.aclose = AsyncMock()
|
||||
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(),
|
||||
sync_stream=False,
|
||||
logging_obj=self._make_logging_obj(),
|
||||
)
|
||||
iterator.async_response_iterator = mock_iterator
|
||||
|
||||
asyncio.run(iterator.aclose())
|
||||
|
||||
mock_iterator.aclose.assert_awaited_once()
|
||||
|
||||
def test_aclose_tolerates_iterator_error(self):
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.aclose = AsyncMock()
|
||||
|
||||
mock_iterator = MagicMock()
|
||||
mock_iterator.aclose = AsyncMock(side_effect=RuntimeError("transport error"))
|
||||
|
||||
iterator = ModelResponseIterator(
|
||||
streaming_response=MagicMock(),
|
||||
sync_stream=False,
|
||||
logging_obj=self._make_logging_obj(),
|
||||
response=mock_response,
|
||||
)
|
||||
iterator.async_response_iterator = mock_iterator
|
||||
|
||||
asyncio.run(iterator.aclose())
|
||||
|
||||
mock_response.aclose.assert_awaited_once()
|
||||
|
||||
def test_custom_stream_wrapper_aclose_triggers_model_response_iterator_aclose(self):
|
||||
"""CustomStreamWrapper.aclose() must propagate to ModelResponseIterator.aclose()."""
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator,
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.aclose = AsyncMock()
|
||||
|
||||
mock_iterator = MagicMock()
|
||||
mock_iterator.aclose = AsyncMock()
|
||||
|
||||
model_response_iter = ModelResponseIterator(
|
||||
streaming_response=MagicMock(),
|
||||
sync_stream=False,
|
||||
logging_obj=self._make_logging_obj(),
|
||||
response=mock_response,
|
||||
)
|
||||
model_response_iter.async_response_iterator = mock_iterator
|
||||
|
||||
wrapper = CustomStreamWrapper(
|
||||
completion_stream=model_response_iter,
|
||||
model="gemini-2.0-flash",
|
||||
custom_llm_provider="vertex_ai",
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
asyncio.run(wrapper.aclose())
|
||||
|
||||
mock_iterator.aclose.assert_awaited_once()
|
||||
mock_response.aclose.assert_awaited_once()
|
||||
|
|
|
|||
|
|
@ -5179,6 +5179,12 @@ async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow():
|
|||
side_effect=lambda update: MCPServer(
|
||||
server_id=legacy_server.server_id,
|
||||
name=legacy_server.name,
|
||||
# Carry alias/server_name forward so get_server_prefix resolves to
|
||||
# "legacy_m2m" (not the server_id) when the request scope filter
|
||||
# matches by alias. Without these, the filter relied on the now-
|
||||
# removed silent fail-open fallback.
|
||||
alias=legacy_server.alias,
|
||||
server_name=legacy_server.server_name,
|
||||
transport=MCPTransport.http,
|
||||
auth_type=legacy_server.auth_type,
|
||||
oauth2_flow=update.get("oauth2_flow", legacy_server.oauth2_flow),
|
||||
|
|
@ -6083,3 +6089,207 @@ async def test_execute_mcp_tool_rest_prefix_retry_resolution_still_enforces_serv
|
|||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail["error"] == "tool_server_mismatch"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression tests for _get_allowed_mcp_servers_from_mcp_server_names
|
||||
#
|
||||
# Prior to the fail-closed fix, an unresolved scope filter (path- or
|
||||
# header-derived) silently returned the caller's full allowed-server set,
|
||||
# which made URL/header namespacing appear to work when it did not.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_mcp_server_for_scope_filter(server_id: str, alias: str) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=alias,
|
||||
alias=alias,
|
||||
server_name=alias,
|
||||
url=f"https://{alias}.test/mcp",
|
||||
transport=MCPTransport.http,
|
||||
mcp_info={"server_name": alias},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_from_mcp_server_names_unknown_name_fails_closed():
|
||||
"""
|
||||
Bug fix: requesting an unknown server name (e.g. ``/mcp/<typo>/``) must
|
||||
NOT silently fall back to the caller's full allowed-server set.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers_from_mcp_server_names,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
allowed = [
|
||||
_make_mcp_server_for_scope_filter("id-a", "alpha"),
|
||||
_make_mcp_server_for_scope_filter("id-b", "beta"),
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp."
|
||||
"MCPRequestHandler._get_mcp_servers_from_access_groups",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[],
|
||||
):
|
||||
result = await _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers=["does-not-exist"],
|
||||
allowed_mcp_servers=allowed,
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_from_mcp_server_names_none_returns_all():
|
||||
"""
|
||||
Regression: ``mcp_servers=None`` (no scope filter requested) must still
|
||||
return the full allowed-server set. This is the legitimate "no scoping"
|
||||
path that the fail-closed fix must not break.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers_from_mcp_server_names,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
allowed = [
|
||||
_make_mcp_server_for_scope_filter("id-a", "alpha"),
|
||||
_make_mcp_server_for_scope_filter("id-b", "beta"),
|
||||
]
|
||||
|
||||
result = await _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers=None,
|
||||
allowed_mcp_servers=allowed,
|
||||
)
|
||||
|
||||
assert {s.server_id for s in result} == {"id-a", "id-b"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_from_mcp_server_names_known_alias_returns_match():
|
||||
"""
|
||||
Regression: a known server alias must still resolve to exactly that
|
||||
server. Guards against the fix accidentally narrowing the happy path.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers_from_mcp_server_names,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
allowed = [
|
||||
_make_mcp_server_for_scope_filter("id-a", "alpha"),
|
||||
_make_mcp_server_for_scope_filter("id-b", "beta"),
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp."
|
||||
"MCPRequestHandler._get_mcp_servers_from_access_groups",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[],
|
||||
):
|
||||
result = await _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers=["alpha"],
|
||||
allowed_mcp_servers=allowed,
|
||||
)
|
||||
|
||||
assert [s.server_id for s in result] == ["id-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_from_mcp_server_names_mixed_known_and_unknown():
|
||||
"""
|
||||
Mixed scope (one valid + one unknown) returns only the resolved server,
|
||||
not the full allowed set. Confirms the fail-closed branch only fires
|
||||
when NOTHING resolves.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers_from_mcp_server_names,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
allowed = [
|
||||
_make_mcp_server_for_scope_filter("id-a", "alpha"),
|
||||
_make_mcp_server_for_scope_filter("id-b", "beta"),
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp."
|
||||
"MCPRequestHandler._get_mcp_servers_from_access_groups",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[],
|
||||
):
|
||||
result = await _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers=["alpha", "does-not-exist"],
|
||||
allowed_mcp_servers=allowed,
|
||||
)
|
||||
|
||||
assert [s.server_id for s in result] == ["id-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_from_mcp_server_names_access_group_resolves():
|
||||
"""
|
||||
Regression: when a requested name is not a server alias but IS an access
|
||||
group, it must still resolve to the underlying servers (not be treated
|
||||
as unresolved).
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers_from_mcp_server_names,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
allowed = [
|
||||
_make_mcp_server_for_scope_filter("id-a", "alpha"),
|
||||
_make_mcp_server_for_scope_filter("id-b", "beta"),
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp."
|
||||
"MCPRequestHandler._get_mcp_servers_from_access_groups",
|
||||
new_callable=AsyncMock,
|
||||
return_value=["id-b"],
|
||||
):
|
||||
result = await _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers=["group-name"],
|
||||
allowed_mcp_servers=allowed,
|
||||
)
|
||||
|
||||
assert [s.server_id for s in result] == ["id-b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_from_mcp_server_names_empty_list_fails_closed():
|
||||
"""
|
||||
Edge case: ``mcp_servers=[]`` (explicit empty scope) is still an
|
||||
explicit filter request. Fail closed rather than returning everything.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_get_allowed_mcp_servers_from_mcp_server_names,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
allowed = [
|
||||
_make_mcp_server_for_scope_filter("id-a", "alpha"),
|
||||
_make_mcp_server_for_scope_filter("id-b", "beta"),
|
||||
]
|
||||
|
||||
result = await _get_allowed_mcp_servers_from_mcp_server_names(
|
||||
mcp_servers=[],
|
||||
allowed_mcp_servers=allowed,
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
|
|
|||
|
|
@ -388,6 +388,60 @@ def test_wildcard_credential_hydration_preserves_deployment_params(
|
|||
}
|
||||
|
||||
|
||||
def test_wildcard_custom_prefix_does_not_stack_provider_prefix(monkeypatch):
|
||||
"""Regression test for #30358.
|
||||
|
||||
A wildcard with a custom prefix (e.g. ``ollama_server1/*`` to distinguish multiple Ollama
|
||||
instances) must not stack the provider's own prefix onto the expanded model ids. The expanded
|
||||
ids should be ``ollama_server1/gemma3:1b`` rather than ``ollama_server1/ollama/gemma3:1b``.
|
||||
"""
|
||||
from litellm.proxy.auth import model_checks
|
||||
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
monkeypatch.setattr(
|
||||
model_checks,
|
||||
"get_provider_models",
|
||||
lambda provider, litellm_params=None: ["ollama/gemma3:1b", "ollama/llama3:8b"],
|
||||
)
|
||||
|
||||
result = get_known_models_from_wildcard(
|
||||
wildcard_model="ollama_server1/*",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="ollama_chat/*", custom_llm_provider="ollama_chat"
|
||||
),
|
||||
)
|
||||
|
||||
assert result == ["ollama_server1/gemma3:1b", "ollama_server1/llama3:8b"]
|
||||
|
||||
|
||||
def test_wildcard_custom_prefix_keeps_org_segment_for_non_provider_first_segment(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Only a known provider prefix should be stripped before re-prefixing.
|
||||
|
||||
If ``get_provider_models`` returns ids whose first segment is an org rather than a litellm
|
||||
provider (e.g. ``meta-llama/Llama-3-8B``), stripping the first slash segment would drop the
|
||||
org and produce an uncallable id. The org segment must be preserved.
|
||||
"""
|
||||
from litellm.proxy.auth import model_checks
|
||||
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
monkeypatch.setattr(
|
||||
model_checks,
|
||||
"get_provider_models",
|
||||
lambda provider, litellm_params=None: ["meta-llama/Llama-3-8B"],
|
||||
)
|
||||
|
||||
result = get_known_models_from_wildcard(
|
||||
wildcard_model="my_hf/*",
|
||||
litellm_params=LiteLLM_Params(model="huggingface/*", custom_llm_provider="huggingface"),
|
||||
)
|
||||
|
||||
assert result == ["my_hf/meta-llama/Llama-3-8B"]
|
||||
|
||||
|
||||
def test_wildcard_credential_hydration_preserves_missing_credential_name(
|
||||
monkeypatch,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -11,7 +11,6 @@ sys.path.insert(
|
|||
|
||||
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -305,3 +304,73 @@ def test_validate_redis_transaction_buffer_passes_when_disabled():
|
|||
general_settings={},
|
||||
redis_usage_cache=None,
|
||||
)
|
||||
|
||||
|
||||
def test_get_transaction_buffer_redis_cache_builds_from_env(monkeypatch):
|
||||
"""
|
||||
When use_redis_transaction_buffer=true, a standalone RedisCache is built from
|
||||
REDIS_* environment variables so the buffer works without a Redis cache backend.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_HOST", "localhost")
|
||||
monkeypatch.setenv("REDIS_PORT", "6379")
|
||||
|
||||
with patch("litellm.proxy.proxy_server.RedisCache") as mock_redis_cache:
|
||||
result = ProxyStartupEvent._get_transaction_buffer_redis_cache(
|
||||
general_settings={"use_redis_transaction_buffer": True},
|
||||
)
|
||||
|
||||
mock_redis_cache.assert_called_once()
|
||||
assert mock_redis_cache.call_args.kwargs["host"] == "localhost"
|
||||
assert result is mock_redis_cache.return_value
|
||||
|
||||
|
||||
def test_get_transaction_buffer_redis_cache_none_when_disabled():
|
||||
"""When use_redis_transaction_buffer is not enabled, no standalone cache is built."""
|
||||
result = ProxyStartupEvent._get_transaction_buffer_redis_cache(
|
||||
general_settings={},
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_get_transaction_buffer_redis_cache_none_without_redis_env():
|
||||
"""
|
||||
When use_redis_transaction_buffer=true but no REDIS_* env vars are set,
|
||||
no standalone cache is built (startup validation then raises the config error).
|
||||
"""
|
||||
with patch("litellm._redis._redis_kwargs_from_environment", return_value={}):
|
||||
result = ProxyStartupEvent._get_transaction_buffer_redis_cache(
|
||||
general_settings={"use_redis_transaction_buffer": True},
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_get_transaction_buffer_redis_cache_none_without_host_or_url():
|
||||
"""
|
||||
A REDIS_* var that is not a connection target (e.g. REDIS_SOCKET_TIMEOUT) must not
|
||||
trigger a build. Without a host or url, get_redis_client raises, so return None and
|
||||
let startup validation surface the config error instead of crashing.
|
||||
"""
|
||||
with patch(
|
||||
"litellm._redis._redis_kwargs_from_environment",
|
||||
return_value={"socket_timeout": 5.0},
|
||||
):
|
||||
result = ProxyStartupEvent._get_transaction_buffer_redis_cache(
|
||||
general_settings={"use_redis_transaction_buffer": True},
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch):
|
||||
"""
|
||||
use_redis_transaction_buffer accepts a string value (e.g. from env/YAML); "true"
|
||||
is parsed to a bool before the standalone cache is built.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_HOST", "localhost")
|
||||
|
||||
with patch("litellm.proxy.proxy_server.RedisCache") as mock_redis_cache:
|
||||
result = ProxyStartupEvent._get_transaction_buffer_redis_cache(
|
||||
general_settings={"use_redis_transaction_buffer": "true"},
|
||||
)
|
||||
|
||||
mock_redis_cache.assert_called_once()
|
||||
assert result is mock_redis_cache.return_value
|
||||
|
|
|
|||
|
|
@ -352,6 +352,79 @@ def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured():
|
|||
assert data["workers"][0]["url"] == "https://worker-1:4001"
|
||||
|
||||
|
||||
def test_ui_discovery_endpoints_hide_default_credentials_hint_default_false():
|
||||
"""Default credentials hint is shown by default (flag false)."""
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
client = TestClient(app)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
|
||||
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None),
|
||||
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False),
|
||||
patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False),
|
||||
):
|
||||
os.environ.pop("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", None)
|
||||
|
||||
response = client.get("/.well-known/litellm-ui-config")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["hide_default_credentials_hint"] is False
|
||||
|
||||
|
||||
def test_ui_discovery_endpoints_hide_default_credentials_hint_via_env_var():
|
||||
"""LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT=true hides the login-page credentials card."""
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
client = TestClient(app)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
|
||||
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None),
|
||||
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False),
|
||||
patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT": "true",
|
||||
"DISABLE_ADMIN_UI": "false",
|
||||
},
|
||||
clear=False,
|
||||
),
|
||||
):
|
||||
|
||||
response = client.get("/.well-known/litellm-ui-config")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["hide_default_credentials_hint"] is True
|
||||
|
||||
|
||||
def test_ui_discovery_endpoints_hide_default_credentials_hint_via_general_settings():
|
||||
"""general_settings.hide_default_credentials_hint=true also hides the card."""
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
client = TestClient(app)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
|
||||
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None),
|
||||
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
{"hide_default_credentials_hint": True},
|
||||
),
|
||||
patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False),
|
||||
):
|
||||
os.environ.pop("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", None)
|
||||
|
||||
response = client.get("/.well-known/litellm-ui-config")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["hide_default_credentials_hint"] is True
|
||||
|
||||
|
||||
def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers():
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
"""
|
||||
Test to verify the Google GenAI proxy API endpoints
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -88,6 +89,8 @@ def test_google_stream_generate_content_endpoint():
|
|||
# stream=True must be forced into the data the processor receives.
|
||||
init_kwargs = mock_init.call_args.kwargs
|
||||
assert init_kwargs["data"]["stream"] is True
|
||||
assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True
|
||||
assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True
|
||||
assert init_kwargs["data"]["model"] == "test-model"
|
||||
assert init_kwargs["data"]["contents"] == [
|
||||
{"role": "user", "parts": [{"text": "Hello"}]}
|
||||
|
|
|
|||
|
|
@ -584,6 +584,50 @@ async def test_logging_hook_multiple_content_items(presidio_guardrail):
|
|||
print("✓ Logging hook multiple content items test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_does_not_mask_pre_call_request(
|
||||
mock_user_api_key, mock_cache
|
||||
):
|
||||
"""
|
||||
A guardrail configured with `logging_only` must only mask PII for logs/traces,
|
||||
never for the request sent to the model. `async_pre_call_hook` should leave the
|
||||
request untouched so the model receives (and replies based on) the real input.
|
||||
|
||||
Regression test for the case where the pre-call hook masked the live request,
|
||||
causing the model's response to contain anonymization tokens (e.g. <PERSON>)
|
||||
instead of the real output.
|
||||
"""
|
||||
presidio_guardrail = _OPTIONAL_PresidioPIIMasking(
|
||||
mock_testing=True,
|
||||
logging_only=True,
|
||||
pii_entities_config={PiiEntityType.PHONE_NUMBER: PiiAction.MASK},
|
||||
)
|
||||
|
||||
async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
|
||||
return text.replace("555-123-4567", "[PHONE]")
|
||||
|
||||
presidio_guardrail.check_pii = mock_check_pii
|
||||
|
||||
original_text = "My phone is 555-123-4567"
|
||||
test_data = {
|
||||
"messages": [{"role": "user", "content": original_text}],
|
||||
"model": "gpt-4",
|
||||
}
|
||||
|
||||
result = await presidio_guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=mock_user_api_key,
|
||||
cache=mock_cache,
|
||||
data=test_data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
# The live request must be unchanged: PII reaches the model intact.
|
||||
assert result["messages"][0]["content"] == original_text
|
||||
assert "[PHONE]" not in result["messages"][0]["content"]
|
||||
|
||||
print("✓ logging_only leaves the pre-call request unmasked")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_presidio_sets_guardrail_information_in_request_data():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ sys.path.insert(
|
|||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
_adjust_dates_for_timezone,
|
||||
_build_aggregated_sql_query,
|
||||
_is_user_agent_tag,
|
||||
get_api_key_metadata,
|
||||
get_daily_activity,
|
||||
|
|
@ -632,6 +634,126 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
|
|||
assert key_data.metrics.spend == 10.0
|
||||
|
||||
|
||||
class TestAdjustDatesForTimezone:
|
||||
"""
|
||||
Regression tests for the timezone double-counting bug.
|
||||
|
||||
Background: the previous implementation expanded the SQL date range by a full
|
||||
UTC day on whichever side a non-UTC timezone offset pointed. Because spend is
|
||||
bucketed in whole UTC days in the aggregation table, that expansion caused
|
||||
single-day queries from non-UTC timezones to include a second full UTC day's
|
||||
worth of data, producing approximately 2x over-counting. The sum of single-day
|
||||
spends across a window then exceeded the equivalent multi-day aggregate, which
|
||||
is mathematically impossible.
|
||||
|
||||
These tests pin the function to a pass-through and assert the additivity
|
||||
invariant that any future implementation must preserve.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"offset_minutes",
|
||||
[
|
||||
None,
|
||||
0,
|
||||
-330, # IST UTC+5:30
|
||||
-540, # JST UTC+9
|
||||
-60, # CET UTC+1
|
||||
240, # AST UTC-4
|
||||
300, # EST UTC-5
|
||||
480, # PST UTC-8
|
||||
],
|
||||
)
|
||||
def test_returns_input_dates_unchanged_for_any_offset(self, offset_minutes):
|
||||
start, end = _adjust_dates_for_timezone(
|
||||
"2026-05-29", "2026-05-29", offset_minutes
|
||||
)
|
||||
assert start == "2026-05-29"
|
||||
assert end == "2026-05-29"
|
||||
|
||||
def test_single_day_query_does_not_widen_to_two_utc_days(self):
|
||||
"""
|
||||
Pins the boundary that caused the original 2x bug: a single IST day must
|
||||
not be translated into a SQL filter covering two UTC days.
|
||||
"""
|
||||
start, end = _adjust_dates_for_timezone("2026-05-29", "2026-05-29", -330)
|
||||
assert start == end == "2026-05-29", (
|
||||
"Single-day IST query expanded to a multi-day UTC range; this is "
|
||||
"the regression that produced approximately 2x over-counting."
|
||||
)
|
||||
|
||||
def test_multi_day_range_endpoints_are_preserved(self):
|
||||
start, end = _adjust_dates_for_timezone("2026-05-29", "2026-06-02", -330)
|
||||
assert (start, end) == ("2026-05-29", "2026-06-02")
|
||||
|
||||
@pytest.mark.parametrize("offset_minutes", [-330, 480])
|
||||
def test_single_day_sums_match_multi_day_window(self, offset_minutes):
|
||||
"""
|
||||
Additivity invariant: querying each day in a window separately and summing
|
||||
the resulting SQL ranges must cover exactly the same range as querying the
|
||||
whole window at once. The bug broke this; without it, single-day sums
|
||||
exceeded the multi-day total by ~50% over a 5-day IST window.
|
||||
"""
|
||||
days = ["2026-05-29", "2026-05-30", "2026-05-31", "2026-06-01", "2026-06-02"]
|
||||
single_day_ranges = [
|
||||
_adjust_dates_for_timezone(d, d, offset_minutes) for d in days
|
||||
]
|
||||
multi_day_range = _adjust_dates_for_timezone(days[0], days[-1], offset_minutes)
|
||||
|
||||
per_day_starts = [r[0] for r in single_day_ranges]
|
||||
per_day_ends = [r[1] for r in single_day_ranges]
|
||||
assert min(per_day_starts) == multi_day_range[0]
|
||||
assert max(per_day_ends) == multi_day_range[1]
|
||||
assert per_day_starts == days
|
||||
assert per_day_ends == days
|
||||
|
||||
|
||||
class TestBuildAggregatedSqlQuery:
|
||||
"""
|
||||
Asserts the SQL emitted by the aggregated query path stays anchored to the
|
||||
user-supplied date range. The original bug shipped a function that returned
|
||||
expanded dates from _adjust_dates_for_timezone, so the regression surface is
|
||||
not just the helper but the SQL it feeds into.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("offset_minutes", [None, 0, -330, 480])
|
||||
def test_sql_date_bounds_are_user_supplied_dates(self, offset_minutes):
|
||||
sql, params = _build_aggregated_sql_query(
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id="user-1",
|
||||
start_date="2026-05-29",
|
||||
end_date="2026-05-29",
|
||||
model=None,
|
||||
api_key=None,
|
||||
timezone_offset_minutes=offset_minutes,
|
||||
)
|
||||
|
||||
assert params[0] == "2026-05-29"
|
||||
assert params[1] == "2026-05-29"
|
||||
assert "date >= $1" in sql
|
||||
assert "date <= $2" in sql
|
||||
|
||||
def test_optional_filters_appear_in_params_in_order(self):
|
||||
sql, params = _build_aggregated_sql_query(
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id="user-1",
|
||||
start_date="2026-05-29",
|
||||
end_date="2026-06-02",
|
||||
model="bedrock/global.anthropic.claude-opus-4-8",
|
||||
api_key="sk-test",
|
||||
timezone_offset_minutes=-330,
|
||||
)
|
||||
|
||||
assert params == [
|
||||
"2026-05-29",
|
||||
"2026-06-02",
|
||||
"user-1",
|
||||
"bedrock/global.anthropic.claude-opus-4-8",
|
||||
"sk-test",
|
||||
]
|
||||
assert "model = $4" in sql
|
||||
assert "api_key = $5" in sql
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_empty_result_set():
|
||||
"""Regression test for the empty-range 500.
|
||||
|
|
|
|||
|
|
@ -11862,7 +11862,6 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(
|
|||
assert "cannot exceed" in msg.lower()
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
|
||||
"""
|
||||
|
|
@ -11941,3 +11940,511 @@ async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
|
|||
assert result["budget_reset_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_key_fn_includes_model_max_budget_usage(monkeypatch):
|
||||
"""
|
||||
/key/info should include model_max_budget_usage showing current-period spend
|
||||
for each model that has a per-model budget configured.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
|
||||
|
||||
test_key_token = "hashed_token_budget_test"
|
||||
model_max_budget = {
|
||||
"gpt-4o": {"budget_limit": 0.50, "time_period": "1d"},
|
||||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.23)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
|
||||
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key_info.token = test_key_token
|
||||
mock_key_info.object_permission_id = None
|
||||
mock_key_info.user_id = "user-x"
|
||||
mock_key_info.team_id = None
|
||||
mock_key_info.litellm_budget_table = None
|
||||
mock_key_info.model_dump.return_value = {
|
||||
"token": test_key_token,
|
||||
"model_max_budget": model_max_budget,
|
||||
"user_id": "user-x",
|
||||
"team_id": None,
|
||||
"object_permission_id": None,
|
||||
"litellm_budget_table": None,
|
||||
}
|
||||
mock_key_info.dict.return_value = mock_key_info.model_dump.return_value
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_key_info
|
||||
)
|
||||
mock_prisma_client.db.query_raw = AsyncMock()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-test-budget-key",
|
||||
)
|
||||
|
||||
result = await info_key_fn(
|
||||
key="sk-test-budget-key",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "model_max_budget_usage" in result["info"]
|
||||
usage = result["info"]["model_max_budget_usage"]
|
||||
assert usage["gpt-4o"]["current_spend"] == 0.23
|
||||
assert usage["gpt-4o"]["budget_limit"] == 0.50
|
||||
assert usage["gpt-4o"]["time_period"] == "1d"
|
||||
mock_prisma_client.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_key_fn_no_model_max_budget_skips_usage(monkeypatch):
|
||||
"""Keys with no model_max_budget should not include model_max_budget_usage."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
|
||||
|
||||
test_key_token = "hashed_token_no_budget"
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
|
||||
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key_info.token = test_key_token
|
||||
mock_key_info.object_permission_id = None
|
||||
mock_key_info.user_id = "user-y"
|
||||
mock_key_info.team_id = None
|
||||
mock_key_info.litellm_budget_table = None
|
||||
mock_key_info.model_dump.return_value = {
|
||||
"token": test_key_token,
|
||||
"model_max_budget": {},
|
||||
"user_id": "user-y",
|
||||
"team_id": None,
|
||||
"object_permission_id": None,
|
||||
"litellm_budget_table": None,
|
||||
}
|
||||
mock_key_info.dict.return_value = mock_key_info.model_dump.return_value
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_key_info
|
||||
)
|
||||
mock_prisma_client.db.query_raw = AsyncMock()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-test-no-budget",
|
||||
)
|
||||
|
||||
result = await info_key_fn(
|
||||
key="sk-test-no-budget",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "model_max_budget_usage" not in result["info"]
|
||||
mock_prisma_client.db.query_raw.assert_not_awaited()
|
||||
mock_user_api_key_cache.async_get_cache.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_key_fn_v2_includes_model_max_budget_usage(monkeypatch):
|
||||
"""/v2/key/info should include model_max_budget_usage for keys with per-model budgets."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
info_key_fn_v2,
|
||||
)
|
||||
|
||||
test_key_token = "hashed_token_v2_test"
|
||||
model_max_budget = {"gpt-4o": {"budget_limit": 1.00, "time_period": "7d"}}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.55)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.token = test_key_token
|
||||
mock_key.user_id = "user-v2"
|
||||
mock_key.team_id = None
|
||||
mock_key.model_dump.return_value = {
|
||||
"token": test_key_token,
|
||||
"model_max_budget": model_max_budget,
|
||||
"user_id": "user-v2",
|
||||
"team_id": None,
|
||||
"litellm_budget_table": None,
|
||||
}
|
||||
mock_key.dict.return_value = mock_key.model_dump.return_value
|
||||
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=[mock_key])
|
||||
mock_prisma_client.db.query_raw = AsyncMock()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
)
|
||||
|
||||
result = await info_key_fn_v2(
|
||||
data=KeyRequest(keys=[test_key_token]),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert len(result["info"]) == 1
|
||||
key_info = result["info"][0]
|
||||
assert "model_max_budget_usage" in key_info
|
||||
usage = key_info["model_max_budget_usage"]
|
||||
assert usage["gpt-4o"]["current_spend"] == 0.55
|
||||
assert usage["gpt-4o"]["budget_limit"] == 1.00
|
||||
assert usage["gpt-4o"]["time_period"] == "7d"
|
||||
mock_prisma_client.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_key_fn_budget_table_fallback(monkeypatch):
|
||||
"""When model_max_budget is empty on the key but set in litellm_budget_table,
|
||||
/key/info should still populate model_max_budget_usage.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
|
||||
|
||||
test_key_token = "hashed_token_budget_table_test"
|
||||
budget_table_model_max_budget = {
|
||||
"bedrock/anthropic.claude-opus-4": {"max_budget": 5, "budget_duration": "30d"},
|
||||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=1.20)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
|
||||
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key_info.token = test_key_token
|
||||
mock_key_info.object_permission_id = None
|
||||
mock_key_info.user_id = "user-bt"
|
||||
mock_key_info.team_id = None
|
||||
mock_key_info.litellm_budget_table = None
|
||||
mock_key_info.model_dump.return_value = {
|
||||
"token": test_key_token,
|
||||
"model_max_budget": {},
|
||||
"user_id": "user-bt",
|
||||
"team_id": None,
|
||||
"object_permission_id": None,
|
||||
"litellm_budget_table": {
|
||||
"budget_id": "bt-123",
|
||||
"budget_duration": "30d",
|
||||
"budget_reset_at": "2026-07-01T00:00:00+00:00",
|
||||
"model_max_budget": budget_table_model_max_budget,
|
||||
},
|
||||
}
|
||||
mock_key_info.dict.return_value = mock_key_info.model_dump.return_value
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_key_info
|
||||
)
|
||||
mock_prisma_client.db.query_raw = AsyncMock()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-test-bt-key",
|
||||
)
|
||||
|
||||
result = await info_key_fn(
|
||||
key="sk-test-bt-key",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "model_max_budget_usage" in result["info"]
|
||||
usage = result["info"]["model_max_budget_usage"]
|
||||
assert usage["bedrock/anthropic.claude-opus-4"]["current_spend"] == 1.20
|
||||
assert usage["bedrock/anthropic.claude-opus-4"]["budget_limit"] == 5
|
||||
assert usage["bedrock/anthropic.claude-opus-4"]["time_period"] == "30d"
|
||||
mock_prisma_client.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_key_fn_v2_budget_table_fallback(monkeypatch):
|
||||
"""When model_max_budget is empty on the key but set in litellm_budget_table,
|
||||
/v2/key/info should still populate model_max_budget_usage."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
info_key_fn_v2,
|
||||
)
|
||||
|
||||
test_key_token = "hashed_token_v2_bt_test"
|
||||
budget_table_model_max_budget = {
|
||||
"bedrock/anthropic.claude-opus-4": {"max_budget": 5, "budget_duration": "30d"},
|
||||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=2.50)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.token = test_key_token
|
||||
mock_key.user_id = "user-v2-bt"
|
||||
mock_key.team_id = None
|
||||
mock_key.model_dump.return_value = {
|
||||
"token": test_key_token,
|
||||
"model_max_budget": {},
|
||||
"user_id": "user-v2-bt",
|
||||
"team_id": None,
|
||||
"litellm_budget_table": {
|
||||
"budget_id": "bt-456",
|
||||
"budget_duration": "30d",
|
||||
"budget_reset_at": "2026-07-01T00:00:00+00:00",
|
||||
"model_max_budget": budget_table_model_max_budget,
|
||||
},
|
||||
}
|
||||
mock_key.dict.return_value = mock_key.model_dump.return_value
|
||||
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=[mock_key])
|
||||
mock_prisma_client.db.query_raw = AsyncMock()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin-v2-bt",
|
||||
)
|
||||
|
||||
result = await info_key_fn_v2(
|
||||
data=KeyRequest(keys=[test_key_token]),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert len(result["info"]) == 1
|
||||
key_info = result["info"][0]
|
||||
assert "model_max_budget_usage" in key_info
|
||||
usage = key_info["model_max_budget_usage"]
|
||||
assert usage["bedrock/anthropic.claude-opus-4"]["current_spend"] == 2.50
|
||||
mock_prisma_client.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_key_fn_provider_prefix_spend_fallback(monkeypatch):
|
||||
"""Cached spend for 'gpt-4o' matches budget key 'openai/gpt-4o' via suffix match."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
|
||||
|
||||
test_key_token = "hashed_token_prefix_test"
|
||||
model_max_budget = {
|
||||
"openai/gpt-4o": {"budget_limit": 2.00, "time_period": "7d"},
|
||||
}
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.75])
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
|
||||
mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key_info.token = test_key_token
|
||||
mock_key_info.object_permission_id = None
|
||||
mock_key_info.user_id = "user-prefix"
|
||||
mock_key_info.team_id = None
|
||||
mock_key_info.litellm_budget_table = None
|
||||
mock_key_info.model_dump.return_value = {
|
||||
"token": test_key_token,
|
||||
"model_max_budget": model_max_budget,
|
||||
"user_id": "user-prefix",
|
||||
"team_id": None,
|
||||
"object_permission_id": None,
|
||||
"litellm_budget_table": None,
|
||||
}
|
||||
mock_key_info.dict.return_value = mock_key_info.model_dump.return_value
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_key_info
|
||||
)
|
||||
mock_prisma_client.db.query_raw = AsyncMock()
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-prefix-test",
|
||||
)
|
||||
|
||||
result = await info_key_fn(
|
||||
key="sk-prefix-test",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "model_max_budget_usage" in result["info"]
|
||||
usage = result["info"]["model_max_budget_usage"]
|
||||
assert usage["openai/gpt-4o"]["current_spend"] == 0.75
|
||||
assert mock_user_api_key_cache.async_get_cache.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_no_cache_returns_empty():
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="some-hash",
|
||||
model_max_budget={"gpt-4o": {"budget_limit": 1.0, "time_period": "1d"}},
|
||||
user_api_key_cache=None,
|
||||
)
|
||||
assert result == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_reads_current_cache_window():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.30)
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="some-hash",
|
||||
model_max_budget={"gpt-4o": {"budget_limit": 1.0, "time_period": "30d"}},
|
||||
user_api_key_cache=mock_user_api_key_cache,
|
||||
)
|
||||
|
||||
assert result["gpt-4o"]["current_spend"] == 0.30
|
||||
mock_user_api_key_cache.async_get_cache.assert_awaited_once_with(
|
||||
key="virtual_key_spend:some-hash:gpt-4o:30d"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_no_duration_in_budget_returns_empty():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock()
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="some-hash",
|
||||
model_max_budget={"gpt-4o": {"budget_limit": 1.0}},
|
||||
user_api_key_cache=mock_user_api_key_cache,
|
||||
)
|
||||
assert result == {}
|
||||
mock_user_api_key_cache.async_get_cache.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_skips_model_without_duration():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.10)
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="some-hash",
|
||||
model_max_budget={
|
||||
"gpt-4o": {"budget_limit": 1.0, "time_period": "1d"},
|
||||
"gpt-3.5-turbo": {"budget_limit": 0.5},
|
||||
},
|
||||
user_api_key_cache=mock_user_api_key_cache,
|
||||
)
|
||||
assert "gpt-4o" in result
|
||||
assert "gpt-3.5-turbo" not in result
|
||||
assert mock_user_api_key_cache.async_get_cache.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_unparseable_duration_skipped():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock()
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="some-hash",
|
||||
model_max_budget={
|
||||
"gpt-4o": {"budget_limit": 1.0, "budget_duration": "not-valid"}
|
||||
},
|
||||
user_api_key_cache=mock_user_api_key_cache,
|
||||
)
|
||||
assert result == {}
|
||||
mock_user_api_key_cache.async_get_cache.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_invalid_budget_config_skipped():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(return_value=0.20)
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="some-hash",
|
||||
model_max_budget={
|
||||
"gpt-4o": {"max_budget": "not-a-number", "budget_duration": "1d"},
|
||||
"gpt-3.5-turbo": {"budget_limit": 0.5, "time_period": "7d"},
|
||||
},
|
||||
user_api_key_cache=mock_user_api_key_cache,
|
||||
)
|
||||
assert "gpt-4o" not in result
|
||||
assert "gpt-3.5-turbo" in result
|
||||
assert mock_user_api_key_cache.async_get_cache.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_model_max_budget_usage_provider_prefix_cache_fallback():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_model_max_budget_usage,
|
||||
)
|
||||
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_user_api_key_cache.async_get_cache = AsyncMock(side_effect=[None, 0.55])
|
||||
|
||||
result = await _build_model_max_budget_usage(
|
||||
api_key_hash="test-hash",
|
||||
model_max_budget={"openai/gpt-4o": {"budget_limit": 2.0, "time_period": "7d"}},
|
||||
user_api_key_cache=mock_user_api_key_cache,
|
||||
)
|
||||
|
||||
assert result["openai/gpt-4o"]["current_spend"] == 0.55
|
||||
assert mock_user_api_key_cache.async_get_cache.await_count == 2
|
||||
|
|
|
|||
|
|
@ -201,6 +201,58 @@ def test_anthropic_provider_fields_support_byok():
|
|||
), "api_base must appear before api_key in credential_fields (matches AI21 and ANTHROPIC_TEXT convention)."
|
||||
|
||||
|
||||
def test_google_ai_studio_provider_fields_expose_api_base():
|
||||
"""The Google AI Studio (gemini) credential form must let admins set a custom
|
||||
api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted
|
||||
proxy at /v1beta) without env var access.
|
||||
|
||||
The runtime gemini provider already supports custom api_base via
|
||||
`vertex_llm_base._check_custom_proxy`; the UI just needs to expose the field.
|
||||
"""
|
||||
app_instance = FastAPI()
|
||||
app_instance.include_router(router)
|
||||
test_client = TestClient(app_instance)
|
||||
|
||||
response = test_client.get("/public/providers/fields")
|
||||
assert response.status_code == 200
|
||||
providers = response.json()
|
||||
|
||||
google_ai = next(
|
||||
(p for p in providers if p["provider"] == "Google_AI_Studio"), None
|
||||
)
|
||||
assert google_ai is not None, "Google_AI_Studio provider entry not found"
|
||||
assert google_ai["litellm_provider"] == "gemini"
|
||||
|
||||
fields_by_key = {f["key"]: f for f in google_ai["credential_fields"]}
|
||||
assert "api_key" in fields_by_key
|
||||
assert "api_base" in fields_by_key, (
|
||||
"Google_AI_Studio provider form must expose api_base so admins can "
|
||||
"point at a Gemini-compatible gateway without env var access."
|
||||
)
|
||||
|
||||
api_base_field = fields_by_key["api_base"]
|
||||
assert api_base_field["required"] is False
|
||||
assert api_base_field["field_type"] == "text"
|
||||
# default_value MUST be null (not the canonical URL): saving it as the
|
||||
# default would persist v1beta into every credential record and bypass
|
||||
# `_get_gemini_url`'s automatic v1alpha routing for Gemini 3+ models. The
|
||||
# placeholder shows the canonical URL so users still get the visual hint.
|
||||
# (See greptileai threads on PR #30419.)
|
||||
assert api_base_field["default_value"] is None
|
||||
assert (
|
||||
api_base_field["placeholder"]
|
||||
== "https://generativelanguage.googleapis.com/v1beta"
|
||||
)
|
||||
|
||||
# UI forms render fields in credential_fields order; api_base should come
|
||||
# first so an admin sees the URL override before the key field (matches
|
||||
# OpenAI and Anthropic conventions).
|
||||
field_order = [f["key"] for f in google_ai["credential_fields"]]
|
||||
assert field_order.index("api_base") < field_order.index(
|
||||
"api_key"
|
||||
), "api_base must appear before api_key in credential_fields."
|
||||
|
||||
|
||||
def test_public_model_hub_with_healthy_model():
|
||||
"""Test that health information is populated for a healthy model"""
|
||||
app = FastAPI()
|
||||
|
|
|
|||
|
|
@ -2737,6 +2737,438 @@ class TestAsyncStreamingDataGeneratorFastPath:
|
|||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
|
||||
class TestDisconnectGatherCleanup:
|
||||
def _disconnect_request(self) -> Request:
|
||||
messages = [
|
||||
{"type": "http.request", "body": b"", "more_body": False},
|
||||
{"type": "http.disconnect"},
|
||||
]
|
||||
|
||||
async def receive():
|
||||
if messages:
|
||||
return messages.pop(0)
|
||||
await asyncio.Event().wait()
|
||||
|
||||
return Request(scope={"type": "http", "headers": []}, receive=receive)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_raises_499_on_client_disconnect(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""With cancel_on_disconnect enabled, base_process_llm_request returns 499."""
|
||||
import asyncio
|
||||
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
async def slow_llm():
|
||||
await asyncio.sleep(9999)
|
||||
|
||||
async def fake_route_request(**_kwargs):
|
||||
return slow_llm()
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.litellm_call_id = "test-call-id"
|
||||
mock_logging_obj._defer_async_logging = False
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.during_call_hook = AsyncMock(return_value=None)
|
||||
mock_proxy_logging._callback_capabilities_cache = {}
|
||||
|
||||
monkeypatch.setattr(cpr, "route_request", fake_route_request)
|
||||
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "gemini-2.0-flash"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await processing_obj.base_process_llm_request(
|
||||
request=self._disconnect_request(),
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
general_settings={"cancel_on_disconnect": True},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
route_type="acompletion",
|
||||
version=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 499
|
||||
assert "disconnected" in exc_info.value.detail.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(
|
||||
self, monkeypatch
|
||||
):
|
||||
import asyncio
|
||||
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
async def fake_gather(*_tasks, **_kwargs):
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.litellm_call_id = "test-call-id"
|
||||
mock_logging_obj._defer_async_logging = False
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.during_call_hook = AsyncMock(return_value=None)
|
||||
mock_proxy_logging._callback_capabilities_cache = {}
|
||||
|
||||
monkeypatch.setattr(cpr.asyncio, "gather", fake_gather)
|
||||
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "gemini-2.0-flash"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
cpr,
|
||||
"route_request",
|
||||
AsyncMock(return_value=asyncio.sleep(9999)),
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await processing_obj.base_process_llm_request(
|
||||
request=mock_request,
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
general_settings={},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
route_type="acompletion",
|
||||
version=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_cancels_during_call_hook_task(self, monkeypatch):
|
||||
import asyncio
|
||||
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
hook_cancelled = False
|
||||
|
||||
async def slow_during_call_hook(**_kwargs):
|
||||
try:
|
||||
await asyncio.sleep(9999)
|
||||
except asyncio.CancelledError:
|
||||
nonlocal hook_cancelled
|
||||
hook_cancelled = True
|
||||
raise
|
||||
|
||||
async def slow_llm():
|
||||
await asyncio.sleep(9999)
|
||||
|
||||
async def fake_route_request(**_kwargs):
|
||||
return slow_llm()
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.litellm_call_id = "test-call-id"
|
||||
mock_logging_obj._defer_async_logging = False
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.during_call_hook = slow_during_call_hook
|
||||
mock_proxy_logging._callback_capabilities_cache = {}
|
||||
|
||||
monkeypatch.setattr(cpr, "route_request", fake_route_request)
|
||||
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "gemini-2.0-flash"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await processing_obj.base_process_llm_request(
|
||||
request=self._disconnect_request(),
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
general_settings={"cancel_on_disconnect": True},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
route_type="acompletion",
|
||||
version=None,
|
||||
)
|
||||
|
||||
assert hook_cancelled is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_pending_gather_tasks_skips_already_done_tasks(self):
|
||||
import asyncio
|
||||
|
||||
from litellm.proxy.common_request_processing import _cancel_pending_gather_tasks
|
||||
|
||||
async def failing_task():
|
||||
raise ValueError("llm api error")
|
||||
|
||||
task = asyncio.create_task(failing_task())
|
||||
with pytest.raises(ValueError, match="llm api error"):
|
||||
await task
|
||||
|
||||
await _cancel_pending_gather_tasks([task])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_pending_gather_tasks_swallows_guardrail_converted_cancel(
|
||||
self,
|
||||
):
|
||||
import asyncio
|
||||
|
||||
from litellm.proxy.common_request_processing import _cancel_pending_gather_tasks
|
||||
|
||||
async def hook_converts_cancel_to_runtime_error():
|
||||
try:
|
||||
await asyncio.sleep(9999)
|
||||
except asyncio.CancelledError:
|
||||
raise RuntimeError("guardrail converted cancel")
|
||||
|
||||
task = asyncio.create_task(hook_converts_cancel_to_runtime_error())
|
||||
await asyncio.sleep(0)
|
||||
await _cancel_pending_gather_tasks([task])
|
||||
assert task.done()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_preserves_llm_error_after_gather(
|
||||
self, monkeypatch
|
||||
):
|
||||
import asyncio
|
||||
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
async def failing_llm():
|
||||
raise ValueError("llm api error")
|
||||
|
||||
async def successful_hook(**_kwargs):
|
||||
return None
|
||||
|
||||
async def fake_route_request(**_kwargs):
|
||||
return failing_llm()
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.litellm_call_id = "test-call-id"
|
||||
mock_logging_obj._defer_async_logging = False
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.during_call_hook = successful_hook
|
||||
mock_proxy_logging._callback_capabilities_cache = {}
|
||||
|
||||
monkeypatch.setattr(cpr, "route_request", fake_route_request)
|
||||
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"model": "gemini-2.0-flash"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=False)
|
||||
mock_request.headers = {}
|
||||
|
||||
with pytest.raises(ValueError, match="llm api error"):
|
||||
await processing_obj.base_process_llm_request(
|
||||
request=mock_request,
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
general_settings={},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
route_type="acompletion",
|
||||
version=None,
|
||||
)
|
||||
|
||||
|
||||
class TestStreamingClientDisconnectLogging:
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_streaming_client_disconnect_sets_error_information(self):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
_record_streaming_client_disconnect_if_needed,
|
||||
)
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {"litellm_params": {}, "metadata": {}}
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=True)
|
||||
request_data = {
|
||||
"litellm_call_id": "test-call-id",
|
||||
"litellm_logging_obj": mock_logging_obj,
|
||||
"metadata": {},
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert (
|
||||
request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
)
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"][
|
||||
"error_information"
|
||||
]["error_code"]
|
||||
== "499"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_streaming_client_disconnect_no_op_when_connected(self):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
_record_streaming_client_disconnect_if_needed,
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=False)
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
|
||||
assert recorded is False
|
||||
assert "client_disconnected" not in request_data["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(
|
||||
self, monkeypatch
|
||||
):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
fire_spy = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.utils.ProxyLogging._fire_deferred_stream_logging",
|
||||
fire_spy,
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=True)
|
||||
mock_response = MagicMock()
|
||||
mock_response.aclose = AsyncMock()
|
||||
request_data = {
|
||||
"metadata": {},
|
||||
"litellm_params": {"metadata": {}},
|
||||
"litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}),
|
||||
}
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
|
||||
request=mock_request,
|
||||
request_data=request_data,
|
||||
response=mock_response,
|
||||
)
|
||||
|
||||
fire_spy.assert_called_once_with(request_data)
|
||||
mock_response.aclose.assert_awaited_once()
|
||||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(
|
||||
self, monkeypatch
|
||||
):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
fire_spy = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.utils.ProxyLogging._fire_deferred_stream_logging",
|
||||
fire_spy,
|
||||
)
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=True)
|
||||
mock_response = MagicMock()
|
||||
mock_response.aclose = AsyncMock()
|
||||
request_data = {"metadata": {}, "litellm_params": {"metadata": {}}}
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
|
||||
request=mock_request,
|
||||
request_data=request_data,
|
||||
response=mock_response,
|
||||
stream_completed=True,
|
||||
)
|
||||
|
||||
fire_spy.assert_not_called()
|
||||
mock_request.is_disconnected.assert_not_awaited()
|
||||
mock_response.aclose.assert_awaited_once()
|
||||
assert "client_disconnected" not in request_data["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_streaming_data_generator_records_499_on_early_aclose(
|
||||
self, monkeypatch
|
||||
):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.utils.ProxyLogging._fire_deferred_stream_logging",
|
||||
MagicMock(),
|
||||
)
|
||||
|
||||
async def mock_streaming_iterator(*_args, **_kwargs):
|
||||
yield {"choices": [{"delta": {"content": "hi"}}]}
|
||||
yield {"choices": [{"delta": {"content": " there"}}]}
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.async_post_call_streaming_iterator_hook = (
|
||||
mock_streaming_iterator
|
||||
)
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=True)
|
||||
mock_response = MagicMock()
|
||||
mock_response.aclose = AsyncMock()
|
||||
request_data = {
|
||||
"model": "gemini-2.0-flash",
|
||||
"metadata": {},
|
||||
"litellm_params": {"metadata": {}},
|
||||
"litellm_logging_obj": MagicMock(
|
||||
model_call_details={"metadata": {}, "litellm_params": {}}
|
||||
),
|
||||
}
|
||||
|
||||
gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=mock_response,
|
||||
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
|
||||
request_data=request_data,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
serialize_chunk=lambda chunk: f"data: {chunk}\n\n",
|
||||
serialize_error=lambda proxy_exc: f"data: {proxy_exc.to_dict()}\n\n",
|
||||
request=mock_request,
|
||||
)
|
||||
await gen.__anext__()
|
||||
await gen.aclose()
|
||||
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
class TestCancelOnDisconnect:
|
||||
"""
|
||||
Coverage for the opt-in `general_settings.cancel_on_disconnect` flag:
|
||||
|
|
|
|||
|
|
@ -188,6 +188,34 @@ async def test_add_litellm_data_to_request_strips_root_pricing_fields():
|
|||
assert "output_cost_per_token" not in updated
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_strips_client_disconnect_metadata():
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"client_disconnected": True,
|
||||
"error_information": {
|
||||
"error_code": "499",
|
||||
"error_message": "Client disconnected the request",
|
||||
"error_class": "ClientDisconnected",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updated = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=_make_request_mock(),
|
||||
user_api_key_dict=_user_api_key_auth(),
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
assert "client_disconnected" not in updated.get("metadata", {})
|
||||
assert "error_information" not in updated.get("metadata", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_strips_metadata_model_info():
|
||||
data = {
|
||||
|
|
|
|||
|
|
@ -5246,10 +5246,10 @@ async def test_async_data_generator_uses_direct_stream_fast_path_without_callbac
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_passes_through_google_native_sse_bytes():
|
||||
async def test_async_data_generator_preserves_non_raw_sse_like_bytes():
|
||||
"""
|
||||
Google-native streamGenerateContent yields raw SSE bytes; they must not be
|
||||
re-wrapped as data: b'data: {...}'.
|
||||
Already formatted SSE bytes from non-raw streams keep the legacy passthrough
|
||||
behavior, including appending a missing event terminator.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
|
|
@ -5305,6 +5305,241 @@ async def test_async_data_generator_passes_through_google_native_sse_bytes():
|
|||
assert yielded_text[-1] == "data: [DONE]\n\n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_buffers_split_google_native_sse_json_frame():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_request_data = {
|
||||
"model": "gemini-3.5-flash",
|
||||
"_litellm_skip_openai_stream_done": True,
|
||||
"_litellm_raw_sse_stream": True,
|
||||
}
|
||||
payload = (
|
||||
'data: {"candidates": [{"content": {"role": "model", "parts": '
|
||||
'[{"text": "", "thoughtSignature": "abc123def456"}]}}]}\n\n'
|
||||
)
|
||||
raw_chunks = [
|
||||
payload[:2].encode("utf-8"),
|
||||
payload[
|
||||
2 : payload.index("thoughtSignature") + len('thoughtSignature": "abc')
|
||||
].encode("utf-8"),
|
||||
payload[
|
||||
payload.index("thoughtSignature") + len('thoughtSignature": "abc') :
|
||||
].encode("utf-8"),
|
||||
]
|
||||
|
||||
class MockStream:
|
||||
def __aiter__(self):
|
||||
return self._stream()
|
||||
|
||||
async def _stream(self):
|
||||
for chunk in raw_chunks:
|
||||
yield chunk
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
mock_response = MockStream()
|
||||
mock_response.aclose = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
|
||||
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
|
||||
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
|
||||
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
|
||||
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj):
|
||||
with patch.object(ProxyLogging, "_fire_deferred_stream_logging"):
|
||||
yielded_data = []
|
||||
async for data in async_data_generator(
|
||||
mock_response, mock_user_api_key_dict, mock_request_data
|
||||
):
|
||||
yielded_data.append(data)
|
||||
|
||||
yielded_text = [
|
||||
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
|
||||
for chunk in yielded_data
|
||||
]
|
||||
|
||||
assert yielded_text == [payload]
|
||||
for chunk in yielded_text:
|
||||
assert chunk.endswith("\n\n")
|
||||
assert json.loads(chunk.removeprefix("data: ").strip())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_flushes_raw_sse_stream_without_trailing_delimiter():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_request_data = {
|
||||
"model": "gemini-3.5-flash",
|
||||
"_litellm_skip_openai_stream_done": True,
|
||||
"_litellm_raw_sse_stream": True,
|
||||
}
|
||||
|
||||
class MockStream:
|
||||
def __aiter__(self):
|
||||
return self._stream()
|
||||
|
||||
async def _stream(self):
|
||||
yield b'data: {"candidates": [{"content": "unterminated"}]'
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
mock_response = MockStream()
|
||||
mock_response.aclose = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
|
||||
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
|
||||
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
|
||||
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
|
||||
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj),
|
||||
patch.object(ProxyLogging, "_fire_deferred_stream_logging"),
|
||||
):
|
||||
yielded_data = []
|
||||
async for data in async_data_generator(
|
||||
mock_response, mock_user_api_key_dict, mock_request_data
|
||||
):
|
||||
yielded_data.append(data)
|
||||
|
||||
yielded_text = [
|
||||
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
|
||||
for chunk in yielded_data
|
||||
]
|
||||
assert len(yielded_text) == 1
|
||||
assert yielded_text[0] == 'data: {"candidates": [{"content": "unterminated"}]\n\n'
|
||||
assert "[DONE]" not in yielded_text[0]
|
||||
mock_proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_errors_when_raw_sse_frame_exceeds_buffer_limit():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_request_data = {
|
||||
"model": "gemini-3.5-flash",
|
||||
"_litellm_skip_openai_stream_done": True,
|
||||
"_litellm_raw_sse_stream": True,
|
||||
}
|
||||
|
||||
class MockStream:
|
||||
def __aiter__(self):
|
||||
return self._stream()
|
||||
|
||||
async def _stream(self):
|
||||
yield b"data: "
|
||||
yield b'{"candidates": [{"content": "unterminated"}]'
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
mock_response = MockStream()
|
||||
mock_response.aclose = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
|
||||
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
|
||||
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
|
||||
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
|
||||
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj),
|
||||
patch("litellm.proxy.proxy_server._MAX_RAW_SSE_BUFFER_CHARS", 8),
|
||||
patch.object(ProxyLogging, "_fire_deferred_stream_logging"),
|
||||
):
|
||||
yielded_data = []
|
||||
async for data in async_data_generator(
|
||||
mock_response, mock_user_api_key_dict, mock_request_data
|
||||
):
|
||||
yielded_data.append(data)
|
||||
|
||||
yielded_text = [
|
||||
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
|
||||
for chunk in yielded_data
|
||||
]
|
||||
assert len(yielded_text) == 1
|
||||
assert "maximum buffered size" in yielded_text[0]
|
||||
assert "[DONE]" not in yielded_text[0]
|
||||
mock_proxy_logging_obj.post_call_failure_hook.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("as_bytes", [True, False])
|
||||
async def test_async_data_generator_checks_raw_sse_buffer_limit_after_complete_frames(
|
||||
as_bytes,
|
||||
):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
complete_frame = 'data: {"candidates": [{"content": "ok"}]}\n\n'
|
||||
partial_frame = "data: "
|
||||
raw_chunk = complete_frame + partial_frame
|
||||
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_request_data = {
|
||||
"model": "gemini-3.5-flash",
|
||||
"_litellm_skip_openai_stream_done": True,
|
||||
"_litellm_raw_sse_stream": True,
|
||||
}
|
||||
|
||||
class MockStream:
|
||||
def __aiter__(self):
|
||||
return self._stream()
|
||||
|
||||
async def _stream(self):
|
||||
yield raw_chunk.encode("utf-8") if as_bytes else raw_chunk
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
mock_response = MockStream()
|
||||
mock_response.aclose = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
|
||||
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
|
||||
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
|
||||
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
|
||||
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj),
|
||||
patch("litellm.proxy.proxy_server._MAX_RAW_SSE_BUFFER_CHARS", 8),
|
||||
patch.object(ProxyLogging, "_fire_deferred_stream_logging"),
|
||||
):
|
||||
yielded_data = []
|
||||
async for data in async_data_generator(
|
||||
mock_response, mock_user_api_key_dict, mock_request_data
|
||||
):
|
||||
yielded_data.append(data)
|
||||
|
||||
yielded_text = [
|
||||
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
|
||||
for chunk in yielded_data
|
||||
]
|
||||
assert yielded_text[0] == complete_frame
|
||||
assert yielded_text[1] == partial_frame + "\n\n"
|
||||
assert "[DONE]" not in "".join(yielded_text)
|
||||
mock_proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_google_genai_stream_omits_openai_done():
|
||||
"""
|
||||
|
|
@ -5359,6 +5594,53 @@ async def test_async_data_generator_google_genai_stream_omits_openai_done():
|
|||
assert "[DONE]" not in "".join(yielded_text)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_does_not_mark_completed_stream_as_disconnect():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.proxy_server import async_data_generator
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_request_data = {"model": "gpt-4o", "metadata": {}}
|
||||
|
||||
class MockStream:
|
||||
def __aiter__(self):
|
||||
return self._stream()
|
||||
|
||||
async def _stream(self):
|
||||
yield {"choices": [{"delta": {"content": "done"}}]}
|
||||
|
||||
async def aclose(self):
|
||||
pass
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.is_disconnected = AsyncMock(return_value=True)
|
||||
mock_response = MockStream()
|
||||
mock_response.aclose = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
|
||||
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
|
||||
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
|
||||
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
|
||||
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
|
||||
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj):
|
||||
with patch.object(ProxyLogging, "_fire_deferred_stream_logging"):
|
||||
yielded_data = []
|
||||
async for data in async_data_generator(
|
||||
mock_response,
|
||||
mock_user_api_key_dict,
|
||||
mock_request_data,
|
||||
request=mock_request,
|
||||
):
|
||||
yielded_data.append(data)
|
||||
|
||||
assert yielded_data[-1] == "data: [DONE]\n\n"
|
||||
mock_request.is_disconnected.assert_not_awaited()
|
||||
assert "client_disconnected" not in mock_request_data["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_google_genai_stream_forwards_error_without_done():
|
||||
"""Stream errors must still reach the client when OpenAI [DONE] is skipped."""
|
||||
|
|
|
|||
|
|
@ -0,0 +1,399 @@
|
|||
"""
|
||||
Unit tests for AWSSecretsManagerV2 cross-region replication via ReplicateSecretToRegions.
|
||||
|
||||
All tests are mocked — no real AWS credentials required.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_CREATE_RESPONSE = {
|
||||
"ARN": "arn:aws:secretsmanager:us-east-1:123456789012:secret:litellm/test-key",
|
||||
"Name": "litellm/test-key",
|
||||
"VersionId": "mock-version-id",
|
||||
}
|
||||
|
||||
_REPLICATE_RESPONSE = {
|
||||
"ARN": "arn:aws:secretsmanager:us-east-1:123456789012:secret:litellm/test-key",
|
||||
"ReplicationStatus": [
|
||||
{"Region": "us-west-2", "Status": "InProgress"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _mock_http_client(json_response: dict) -> MagicMock:
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
mock_response.json.return_value = json_response
|
||||
mock_async_client = AsyncMock()
|
||||
mock_async_client.post.return_value = mock_response
|
||||
return mock_async_client
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: async_write_secret + replication
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_replicates_when_configured():
|
||||
"""async_replicate_secret is called after a successful CreateSecret when replica_regions is set."""
|
||||
manager = AWSSecretsManagerV2(replica_regions=["us-west-2"])
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b'{"Name":"litellm/test-key"}',
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client(_CREATE_RESPONSE),
|
||||
):
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"async_replicate_secret",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_REPLICATE_RESPONSE,
|
||||
) as mock_replicate:
|
||||
result = await manager.async_write_secret(
|
||||
secret_name="litellm/test-key",
|
||||
secret_value="sk-test-value",
|
||||
)
|
||||
|
||||
assert result == _CREATE_RESPONSE
|
||||
mock_replicate.assert_called_once_with(
|
||||
secret_name="litellm/test-key",
|
||||
replica_regions=["us-west-2"],
|
||||
optional_params=None,
|
||||
timeout=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_no_replication_when_not_configured():
|
||||
"""async_replicate_secret is NOT called when replica_regions is None."""
|
||||
manager = AWSSecretsManagerV2(replica_regions=None)
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b'{"Name":"litellm/test-key"}',
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client(_CREATE_RESPONSE),
|
||||
):
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"async_replicate_secret",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_replicate:
|
||||
result = await manager.async_write_secret(
|
||||
secret_name="litellm/test-key",
|
||||
secret_value="sk-test-value",
|
||||
)
|
||||
|
||||
assert result == _CREATE_RESPONSE
|
||||
mock_replicate.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replication_failure_does_not_fail_write():
|
||||
"""If async_replicate_secret raises, async_write_secret still returns the CreateSecret response."""
|
||||
manager = AWSSecretsManagerV2(replica_regions=["us-west-2"])
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b'{"Name":"litellm/test-key"}',
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client(_CREATE_RESPONSE),
|
||||
):
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"async_replicate_secret",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=ValueError("AccessDenied: not authorized"),
|
||||
):
|
||||
result = await manager.async_write_secret(
|
||||
secret_name="litellm/test-key",
|
||||
secret_value="sk-test-value",
|
||||
)
|
||||
|
||||
assert result == _CREATE_RESPONSE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: async_replicate_secret directly
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_replicate_secret_empty_regions_returns_empty():
|
||||
"""async_replicate_secret returns {} immediately for an empty list — no HTTP call."""
|
||||
manager = AWSSecretsManagerV2()
|
||||
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
result = await manager.async_replicate_secret(
|
||||
secret_name="litellm/test-key",
|
||||
replica_regions=[],
|
||||
)
|
||||
|
||||
assert result == {}
|
||||
mock_get_client.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_replicate_secret_correct_payload():
|
||||
"""async_replicate_secret sends the correct AddReplicaRegions payload."""
|
||||
manager = AWSSecretsManagerV2()
|
||||
captured: dict = {}
|
||||
|
||||
def capture_prepare(action, secret_name, optional_params=None, request_data=None):
|
||||
captured.update(request_data or {})
|
||||
captured["_action"] = action
|
||||
return (
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b"{}",
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2, "_prepare_request", side_effect=capture_prepare
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client(_REPLICATE_RESPONSE),
|
||||
):
|
||||
result = await manager.async_replicate_secret(
|
||||
secret_name="litellm/test-key",
|
||||
replica_regions=["us-west-2", "eu-west-1"],
|
||||
)
|
||||
|
||||
assert result == _REPLICATE_RESPONSE
|
||||
assert captured["_action"] == "ReplicateSecretToRegions"
|
||||
assert captured["SecretId"] == "litellm/test-key"
|
||||
assert captured["AddReplicaRegions"] == [
|
||||
{"Region": "us-west-2"},
|
||||
{"Region": "eu-west-1"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replication_fires_on_create(caplog):
|
||||
"""async_replicate_secret emits an INFO log line mentioning ReplicateSecretToRegions."""
|
||||
import logging
|
||||
|
||||
manager = AWSSecretsManagerV2()
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b"{}",
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client(_REPLICATE_RESPONSE),
|
||||
):
|
||||
with caplog.at_level(logging.INFO, logger="LiteLLM"):
|
||||
await manager.async_replicate_secret(
|
||||
secret_name="litellm/test-key",
|
||||
replica_regions=["us-west-2"],
|
||||
)
|
||||
|
||||
assert "ReplicateSecretToRegions" in caplog.text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: load_aws_secret_manager forwards replica_regions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_load_aws_secret_manager_passes_replica_regions():
|
||||
"""load_aws_secret_manager must forward replica_regions from key_management_settings."""
|
||||
import litellm
|
||||
|
||||
original = litellm.secret_manager_client
|
||||
settings = MagicMock()
|
||||
settings.aws_region_name = "us-east-1"
|
||||
settings.aws_role_name = None
|
||||
settings.aws_session_name = None
|
||||
settings.aws_external_id = None
|
||||
settings.aws_profile_name = None
|
||||
settings.aws_web_identity_token = None
|
||||
settings.aws_sts_endpoint = None
|
||||
settings.replica_regions = ["us-west-2", "eu-west-1"]
|
||||
|
||||
try:
|
||||
AWSSecretsManagerV2.load_aws_secret_manager(
|
||||
use_aws_secret_manager=True,
|
||||
key_management_settings=settings,
|
||||
)
|
||||
|
||||
assert isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
|
||||
assert litellm.secret_manager_client.replica_regions == [
|
||||
"us-west-2",
|
||||
"eu-west-1",
|
||||
]
|
||||
finally:
|
||||
litellm.secret_manager_client = original
|
||||
|
||||
|
||||
def _http_status_error(status_code: int, body: str) -> httpx.HTTPStatusError:
|
||||
request = httpx.Request("POST", "https://secretsmanager.us-east-1.amazonaws.com")
|
||||
response = httpx.Response(status_code=status_code, text=body, request=request)
|
||||
return httpx.HTTPStatusError(message=body, request=request, response=response)
|
||||
|
||||
|
||||
def _mock_http_client_raising(exc: Exception) -> MagicMock:
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status.side_effect = exc
|
||||
mock_async_client = AsyncMock()
|
||||
mock_async_client.post.return_value = mock_response
|
||||
return mock_async_client
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: error paths in async_write_secret
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_http_error_raises():
|
||||
"""async_write_secret raises ValueError when CreateSecret returns a non-2xx HTTP status."""
|
||||
manager = AWSSecretsManagerV2()
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b'{"Name":"litellm/test-key"}',
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client_raising(
|
||||
_http_status_error(400, "ResourceExistsException")
|
||||
),
|
||||
):
|
||||
with pytest.raises(ValueError, match="HTTP error occurred"):
|
||||
await manager.async_write_secret(
|
||||
secret_name="litellm/test-key",
|
||||
secret_value="sk-test-value",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_secret_timeout_raises():
|
||||
"""async_write_secret raises ValueError when the CreateSecret call times out."""
|
||||
manager = AWSSecretsManagerV2()
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b'{"Name":"litellm/test-key"}',
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client_raising(
|
||||
httpx.ReadTimeout("timed out", request=None)
|
||||
),
|
||||
):
|
||||
with pytest.raises(ValueError, match="Timeout error occurred"):
|
||||
await manager.async_write_secret(
|
||||
secret_name="litellm/test-key",
|
||||
secret_value="sk-test-value",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: error paths in async_replicate_secret
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replicate_secret_http_error_raises():
|
||||
"""async_replicate_secret raises ValueError when ReplicateSecretToRegions returns a non-2xx status."""
|
||||
manager = AWSSecretsManagerV2()
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b"{}",
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client_raising(
|
||||
_http_status_error(403, "AccessDeniedException")
|
||||
),
|
||||
):
|
||||
with pytest.raises(ValueError, match="HTTP error occurred"):
|
||||
await manager.async_replicate_secret(
|
||||
secret_name="litellm/test-key",
|
||||
replica_regions=["us-west-2"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replicate_secret_timeout_raises():
|
||||
"""async_replicate_secret raises ValueError when the ReplicateSecretToRegions call times out."""
|
||||
manager = AWSSecretsManagerV2()
|
||||
|
||||
with patch.object(
|
||||
AWSSecretsManagerV2,
|
||||
"_prepare_request",
|
||||
return_value=(
|
||||
"https://secretsmanager.us-east-1.amazonaws.com",
|
||||
{"Content-Type": "application/x-amz-json-1.1"},
|
||||
b"{}",
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"litellm.secret_managers.aws_secret_manager_v2.get_async_httpx_client",
|
||||
return_value=_mock_http_client_raising(
|
||||
httpx.ReadTimeout("timed out", request=None)
|
||||
),
|
||||
):
|
||||
with pytest.raises(ValueError, match="Timeout error occurred"):
|
||||
await manager.async_replicate_secret(
|
||||
secret_name="litellm/test-key",
|
||||
replica_regions=["us-west-2"],
|
||||
)
|
||||
|
|
@ -0,0 +1,89 @@
|
|||
"""
|
||||
Validate that the native (first-party) Anthropic Claude Sonnet 4.5 / 4.6 entries
|
||||
carry the 1-hour prompt-cache write tier (`cache_creation_input_token_cost_above_1hr`)
|
||||
in `model_prices_and_context_window.json`.
|
||||
|
||||
Anthropic's first-party API charges a separate 1-hour cache write rate (2x base
|
||||
input) alongside the 5-minute write (1.25x base input) and cache read (0.1x base
|
||||
input). The 1h/5m ratio is therefore 1.6. Without the 1-hour field, cost tracking
|
||||
on 1-hour-TTL prompt caching falls back to the 5-minute rate and undercounts spend.
|
||||
|
||||
The native (non-bedrock) `claude-sonnet-4-5*` / `claude-sonnet-4-6` entries were
|
||||
missing this field, while every sibling (`vertex_ai/`, `azure_ai/`, the
|
||||
`*.anthropic.*` Bedrock profiles) and the older `claude-sonnet-4-20250514` already
|
||||
carried it. This test guards against regression.
|
||||
|
||||
Values (per token):
|
||||
Sonnet base input 3e-06 -> 5m 3.75e-06, 1h 6e-06
|
||||
Sonnet 4.5 long-context (>200K) base 6e-06 -> 5m 7.5e-06, 1h 1.2e-05
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_data():
|
||||
json_path = os.path.join(
|
||||
os.path.dirname(__file__), "../../model_prices_and_context_window.json"
|
||||
)
|
||||
with open(json_path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
# (model_key, expected 1hr write per token, expected 1hr long-context tier or None)
|
||||
EXPECTED = [
|
||||
("claude-sonnet-4-5", 6e-06, 1.2e-05),
|
||||
("claude-sonnet-4-5-20250929", 6e-06, 1.2e-05),
|
||||
("claude-sonnet-4-5-20250929-v1:0", 6e-06, 1.2e-05),
|
||||
("claude-sonnet-4-6", 6e-06, None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_key, expected_1hr, expected_1hr_lc", EXPECTED)
|
||||
def test_anthropic_sonnet_1hr_cache_write_pricing(
|
||||
model_data, model_key, expected_1hr, expected_1hr_lc
|
||||
):
|
||||
assert model_key in model_data, f"Missing model entry: {model_key}"
|
||||
info = model_data[model_key]
|
||||
|
||||
# Regular 1hr cache write rate must be present and exact.
|
||||
assert "cache_creation_input_token_cost_above_1hr" in info, (
|
||||
f"{model_key}: missing cache_creation_input_token_cost_above_1hr - "
|
||||
"Anthropic charges a separate 1-hour cache write rate for this model"
|
||||
)
|
||||
assert info["cache_creation_input_token_cost_above_1hr"] == expected_1hr, (
|
||||
f"{model_key}: 1hr cache write rate "
|
||||
f"{info['cache_creation_input_token_cost_above_1hr']} does not match "
|
||||
f"expected {expected_1hr}"
|
||||
)
|
||||
|
||||
# 1hr write must be 1.6x the 5-minute write (Anthropic 2x-base / 1.25x-base).
|
||||
ratio = (
|
||||
info["cache_creation_input_token_cost_above_1hr"]
|
||||
/ info["cache_creation_input_token_cost"]
|
||||
)
|
||||
assert (
|
||||
abs(ratio - 1.6) < 1e-9
|
||||
), f"{model_key}: 1hr/5min ratio is {ratio}, expected 1.6"
|
||||
|
||||
# Long-context (>200K) 1hr tier, where the model publishes a >200K tier.
|
||||
if expected_1hr_lc is not None:
|
||||
assert (
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens" in info
|
||||
), f"{model_key}: missing 1hr cache write tier for >200K context"
|
||||
assert (
|
||||
info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"]
|
||||
== expected_1hr_lc
|
||||
)
|
||||
ratio_lc = (
|
||||
info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"]
|
||||
/ info["cache_creation_input_token_cost_above_200k_tokens"]
|
||||
)
|
||||
assert (
|
||||
abs(ratio_lc - 1.6) < 1e-9
|
||||
), f"{model_key}: long-context 1hr/5min ratio is {ratio_lc}, expected 1.6"
|
||||
else:
|
||||
assert "cache_creation_input_token_cost_above_1hr_above_200k_tokens" not in info
|
||||
68
tests/test_litellm/test_gpt_5_5_model_metadata.py
Normal file
68
tests/test_litellm/test_gpt_5_5_model_metadata.py
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"])
|
||||
def test_azure_ai_gpt_5_5_model_info(model):
|
||||
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
|
||||
with open(json_path) as f:
|
||||
model_cost = json.load(f)
|
||||
|
||||
info = model_cost.get(model)
|
||||
assert (
|
||||
info is not None
|
||||
), f"{model} not found in model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "azure_ai"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == 5e-06
|
||||
assert info["output_cost_per_token"] == 3e-05
|
||||
assert info["cache_read_input_token_cost"] == 5e-07
|
||||
|
||||
assert info["input_cost_per_token_above_272k_tokens"] == 1e-05
|
||||
assert info["output_cost_per_token_above_272k_tokens"] == 4.5e-05
|
||||
assert info["cache_read_input_token_cost_above_272k_tokens"] == 1e-06
|
||||
|
||||
assert info["input_cost_per_token_priority"] == 1e-05
|
||||
assert info["output_cost_per_token_priority"] == 6e-05
|
||||
|
||||
assert info["max_input_tokens"] == 1050000
|
||||
assert info["max_output_tokens"] == 128000
|
||||
assert info["max_tokens"] == 128000
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_prompt_caching"] is True
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_web_search"] is True
|
||||
# gpt-5.5 dropped minimal reasoning effort support (true on gpt-5.4)
|
||||
assert info["supports_minimal_reasoning_effort"] is False
|
||||
|
||||
routed_model, provider, _, _ = get_llm_provider(model=model)
|
||||
assert routed_model == model.split("/", 1)[1]
|
||||
# azure_ai/* models resolve under the azure provider in get_llm_provider
|
||||
assert provider == "azure"
|
||||
|
||||
|
||||
def test_azure_ai_gpt_5_5_backup_matches_main():
|
||||
"""Ensure the bundled model cost map stays in sync with the canonical file."""
|
||||
repo_root = Path(__file__).parents[2]
|
||||
main_path = repo_root / "model_prices_and_context_window.json"
|
||||
backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json"
|
||||
|
||||
with open(main_path) as f:
|
||||
main_cost = json.load(f)
|
||||
with open(backup_path) as f:
|
||||
backup_cost = json.load(f)
|
||||
|
||||
for model in ("azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"):
|
||||
assert backup_cost.get(model) == main_cost.get(
|
||||
model
|
||||
), f"{model} differs between main and backup model cost maps"
|
||||
55
tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
Normal file
55
tests/test_litellm/test_mistral_medium_3_5_model_metadata.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["mistral/mistral-medium-3-5"])
|
||||
def test_mistral_medium_3_5_model_info(model):
|
||||
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
|
||||
with open(json_path) as f:
|
||||
model_cost = json.load(f)
|
||||
|
||||
info = model_cost.get(model)
|
||||
assert (
|
||||
info is not None
|
||||
), f"{model} not found in model_prices_and_context_window.json"
|
||||
|
||||
assert info["litellm_provider"] == "mistral"
|
||||
assert info["mode"] == "chat"
|
||||
|
||||
assert info["input_cost_per_token"] == 1.5e-06
|
||||
assert info["output_cost_per_token"] == 7.5e-06
|
||||
|
||||
assert info["max_input_tokens"] == 262144
|
||||
assert info["max_output_tokens"] == 262144
|
||||
assert info["max_tokens"] == 262144
|
||||
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_response_schema"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
assert info["supports_vision"] is True
|
||||
assert info["supports_assistant_prefill"] is True
|
||||
|
||||
routed_model, provider, _, _ = get_llm_provider(model=model)
|
||||
assert routed_model == model.split("/", 1)[1]
|
||||
assert provider == "mistral"
|
||||
|
||||
|
||||
def test_mistral_medium_3_5_backup_matches_main():
|
||||
"""Ensure the bundled model cost map stays in sync with the canonical file."""
|
||||
repo_root = Path(__file__).parents[2]
|
||||
main_path = repo_root / "model_prices_and_context_window.json"
|
||||
backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json"
|
||||
|
||||
with open(main_path) as f:
|
||||
main_cost = json.load(f)
|
||||
with open(backup_path) as f:
|
||||
backup_cost = json.load(f)
|
||||
|
||||
for model in ("mistral/mistral-medium-3-5",):
|
||||
assert backup_cost.get(model) == main_cost.get(
|
||||
model
|
||||
), f"{model} differs between main and backup model cost maps"
|
||||
|
|
@ -365,3 +365,37 @@ async def test_router_order_fallback_with_wildcard_model_group():
|
|||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
assert response._hidden_params["model_id"] == "2"
|
||||
|
||||
|
||||
def test_check_non_standard_fallback_format():
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
_check_non_standard_fallback_format,
|
||||
)
|
||||
|
||||
# Standard formats
|
||||
assert (
|
||||
_check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}])
|
||||
== False
|
||||
)
|
||||
assert _check_non_standard_fallback_format([{"model": ["qwen-backup"]}]) == False
|
||||
assert (
|
||||
_check_non_standard_fallback_format(
|
||||
[{"model": ["qwen-backup"], "region": ["us-east-1"]}]
|
||||
)
|
||||
== False
|
||||
)
|
||||
|
||||
# Non-standard formats
|
||||
assert _check_non_standard_fallback_format([{"model": "qwen-backup"}]) == True
|
||||
assert (
|
||||
_check_non_standard_fallback_format(
|
||||
[{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}]
|
||||
)
|
||||
== True
|
||||
)
|
||||
assert (
|
||||
_check_non_standard_fallback_format(
|
||||
[{"model": ["qwen-backup"], "api_key": "some-key"}]
|
||||
)
|
||||
== True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -890,6 +890,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"/v1/audio/speech",
|
||||
"/v1/ocr",
|
||||
"/vertex_ai/live",
|
||||
"/v1/realtime/transcription_sessions",
|
||||
],
|
||||
},
|
||||
},
|
||||
|
|
@ -4153,6 +4154,96 @@ class TestValidateAndFixThinkingParam:
|
|||
assert "budget_tokens" not in thinking
|
||||
|
||||
|
||||
def test_deepseek_v4_models_in_cost_map():
|
||||
"""
|
||||
Test that deepseek-v4-flash and deepseek-v4-pro entries are correctly
|
||||
configured in model_prices_and_context_window.json.
|
||||
|
||||
Prices sourced from https://api-docs.deepseek.com/quick_start/pricing:
|
||||
- deepseek-v4-flash: $0.14/M input, $0.28/M output
|
||||
- deepseek-v4-pro: $0.435/M input, $0.87/M output (75% discounted active price)
|
||||
|
||||
Closes https://github.com/BerriAI/litellm/issues/26709
|
||||
"""
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
|
||||
with open(json_path) as f:
|
||||
model_cost = json.load(f)
|
||||
|
||||
# --- bare model names ---
|
||||
for key, expected_input, expected_output, expected_cache in [
|
||||
("deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09),
|
||||
("deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09),
|
||||
]:
|
||||
info = model_cost.get(key)
|
||||
assert info is not None, f"{key} missing from model_prices_and_context_window.json"
|
||||
assert info["litellm_provider"] == "deepseek"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
assert info["cache_read_input_token_cost"] == expected_cache
|
||||
assert info["max_input_tokens"] == 1_000_000
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
|
||||
# --- provider-prefixed names ---
|
||||
for key, expected_input, expected_output, expected_cache in [
|
||||
("deepseek/deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09),
|
||||
("deepseek/deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09),
|
||||
]:
|
||||
info = model_cost.get(key)
|
||||
assert info is not None, f"{key} missing from model_prices_and_context_window.json"
|
||||
assert info["litellm_provider"] == "deepseek"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
assert info["cache_read_input_token_cost"] == expected_cache
|
||||
assert info["supports_function_calling"] is True
|
||||
assert info["supports_tool_choice"] is True
|
||||
|
||||
|
||||
def test_deepseek_v4_models_in_backup_cost_map():
|
||||
"""
|
||||
Test that deepseek-v4-flash and deepseek-v4-pro entries are correctly
|
||||
configured in litellm/model_prices_and_context_window_backup.json.
|
||||
"""
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
json_path = Path(__file__).parents[2] / "litellm" / "model_prices_and_context_window_backup.json"
|
||||
with open(json_path) as f:
|
||||
model_cost = json.load(f)
|
||||
|
||||
# --- bare model names ---
|
||||
for key, expected_input, expected_output, expected_cache in [
|
||||
("deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09),
|
||||
("deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09),
|
||||
]:
|
||||
info = model_cost.get(key)
|
||||
assert info is not None, f"{key} missing from backup JSON"
|
||||
assert info["litellm_provider"] == "deepseek"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
assert info["cache_read_input_token_cost"] == expected_cache
|
||||
assert info["max_input_tokens"] == 1_000_000
|
||||
|
||||
# --- provider-prefixed names ---
|
||||
for key, expected_input, expected_output, expected_cache in [
|
||||
("deepseek/deepseek-v4-flash", 1.4e-07, 2.8e-07, 2.8e-09),
|
||||
("deepseek/deepseek-v4-pro", 4.35e-07, 8.7e-07, 3.625e-09),
|
||||
]:
|
||||
info = model_cost.get(key)
|
||||
assert info is not None, f"{key} missing from backup JSON"
|
||||
assert info["litellm_provider"] == "deepseek"
|
||||
assert info["mode"] == "chat"
|
||||
assert info["input_cost_per_token"] == expected_input
|
||||
assert info["output_cost_per_token"] == expected_output
|
||||
assert info["cache_read_input_token_cost"] == expected_cache
|
||||
|
||||
|
||||
class TestBedrockBaseModelLabelKeepsTools:
|
||||
"""Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly
|
||||
label must not silently drop ``tools``/``tool_choice`` under ``drop_params``."""
|
||||
|
|
@ -4217,3 +4308,4 @@ def test_aws_bedrock_project_id_excluded_from_bedrock_optional_params():
|
|||
|
||||
assert "aws_bedrock_project_id" not in result
|
||||
assert result["aws_region_name"] == "us-east-1"
|
||||
|
||||
|
|
|
|||
|
|
@ -188,28 +188,30 @@ function LoginPageContent() {
|
|||
<Text type="secondary">Access your LiteLLM Admin UI.</Text>
|
||||
</div>
|
||||
|
||||
<Alert
|
||||
message="Default Credentials"
|
||||
description={
|
||||
<>
|
||||
<Paragraph className="text-sm">
|
||||
By default, Username is <code className="bg-gray-100 px-1 py-0.5 rounded text-xs">admin</code> and
|
||||
Password is your set LiteLLM Proxy
|
||||
<code className="bg-gray-100 px-1 py-0.5 rounded text-xs">MASTER_KEY</code>.
|
||||
</Paragraph>
|
||||
<Paragraph className="text-sm">
|
||||
Need to set UI credentials or SSO?{" "}
|
||||
<a href="https://docs.litellm.ai/docs/proxy/ui" target="_blank" rel="noopener noreferrer">
|
||||
Check the documentation
|
||||
</a>
|
||||
.
|
||||
</Paragraph>
|
||||
</>
|
||||
}
|
||||
type="info"
|
||||
icon={<InfoCircleOutlined />}
|
||||
showIcon
|
||||
/>
|
||||
{!uiConfig?.hide_default_credentials_hint && (
|
||||
<Alert
|
||||
message="Default Credentials"
|
||||
description={
|
||||
<>
|
||||
<Paragraph className="text-sm">
|
||||
By default, Username is <code className="bg-gray-100 px-1 py-0.5 rounded text-xs">admin</code> and
|
||||
Password is your set LiteLLM Proxy
|
||||
<code className="bg-gray-100 px-1 py-0.5 rounded text-xs">MASTER_KEY</code>.
|
||||
</Paragraph>
|
||||
<Paragraph className="text-sm">
|
||||
Need to set UI credentials or SSO?{" "}
|
||||
<a href="https://docs.litellm.ai/docs/proxy/ui" target="_blank" rel="noopener noreferrer">
|
||||
Check the documentation
|
||||
</a>
|
||||
.
|
||||
</Paragraph>
|
||||
</>
|
||||
}
|
||||
type="info"
|
||||
icon={<InfoCircleOutlined />}
|
||||
showIcon
|
||||
/>
|
||||
)}
|
||||
|
||||
{error && <Alert message={error} type="error" showIcon />}
|
||||
|
||||
|
|
|
|||
|
|
@ -55,4 +55,39 @@ describe("LoggingCallbacksTable", () => {
|
|||
);
|
||||
expect(getByText("custom_callback_x")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Regression: `/get_callbacks` returns the same `name` twice when a
|
||||
// callback is registered for both success and failure (e.g. `generic_api`
|
||||
// → POST to spend-log on both 200 and 4xx/5xx). The UI used to ignore
|
||||
// the `type` field and render every row as "Success", masking the
|
||||
// failure registration. Reading `record.type` fixes the badge AND
|
||||
// composing the rowKey with type avoids React's duplicate-key warning.
|
||||
it("renders distinct Success and Failure badges for same-name dual registration", () => {
|
||||
const baseVars = {
|
||||
SLACK_WEBHOOK_URL: null,
|
||||
LANGFUSE_PUBLIC_KEY: null,
|
||||
LANGFUSE_SECRET_KEY: null,
|
||||
LANGFUSE_HOST: null,
|
||||
OPENMETER_API_KEY: null,
|
||||
};
|
||||
const { getAllByText, getByText } = render(
|
||||
<LoggingCallbacksTable
|
||||
callbacks={[
|
||||
{ name: "generic_api", type: "success", variables: baseVars },
|
||||
{ name: "generic_api", type: "failure", variables: baseVars },
|
||||
]}
|
||||
availableCallbacks={{
|
||||
generic_api: {
|
||||
litellm_callback_name: "generic_api",
|
||||
litellm_callback_params: [],
|
||||
ui_callback_name: "Custom Callback API",
|
||||
},
|
||||
}}
|
||||
/>,
|
||||
);
|
||||
// Both rows show the same display name, but distinct mode badges.
|
||||
expect(getAllByText("Custom Callback API")).toHaveLength(2);
|
||||
expect(getByText("Success")).toBeInTheDocument();
|
||||
expect(getByText("Failure")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -48,7 +48,6 @@ export const LoggingCallbacksTable: React.FC<LoggingCallbacksProps> = ({
|
|||
key: "name",
|
||||
render: (_: string, record: CallbackRow) => {
|
||||
const id = record.name;
|
||||
console.log("availableCallbacks", availableCallbacks);
|
||||
const displayName = availableCallbacks[id]?.ui_callback_name || id;
|
||||
return <div className="font-medium text-gray-800">{displayName}</div>;
|
||||
},
|
||||
|
|
@ -57,7 +56,10 @@ export const LoggingCallbacksTable: React.FC<LoggingCallbacksProps> = ({
|
|||
title: <span className="font-medium text-gray-700">Mode</span>,
|
||||
key: "mode",
|
||||
render: (_: unknown, record: CallbackRow) => {
|
||||
const mode = record.mode || "success";
|
||||
// Backend sends `type` (success | failure); legacy in-memory rows
|
||||
// from add-callback flow set `mode`. Read both so newly-added rows
|
||||
// and server-fetched rows both render correctly.
|
||||
const mode = record.type || record.mode || "success";
|
||||
const label = CALLBACK_MODES.find((m) => m.value === mode)?.label || mode;
|
||||
const badgeClass =
|
||||
mode === "success"
|
||||
|
|
@ -109,7 +111,10 @@ export const LoggingCallbacksTable: React.FC<LoggingCallbacksProps> = ({
|
|||
<Table
|
||||
columns={columns}
|
||||
dataSource={callbacks as CallbackRow[]}
|
||||
rowKey={(record) => record.name}
|
||||
// `generic_api` can appear as both a success and a failure
|
||||
// callback simultaneously — keying by `name` alone produced
|
||||
// duplicate React keys. Compose with type to keep keys unique.
|
||||
rowKey={(record) => `${record.name}-${record.type || record.mode || "success"}`}
|
||||
pagination={false}
|
||||
rowClassName={() => "hover:bg-gray-50"}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -1,5 +1,12 @@
|
|||
export interface AlertingObject {
|
||||
name: string;
|
||||
// Backend distinguishes success vs failure callback registrations
|
||||
// (`/get_callbacks` returns `type: "success" | "failure"`). Same callback
|
||||
// (e.g. `generic_api`) can appear twice — once per event class — and
|
||||
// those entries fire on disjoint events, not double-fire on one event.
|
||||
// UI must read this to render the correct badge; missing it caused
|
||||
// every row to render as "Success".
|
||||
type?: "success" | "failure" | "success_and_failure";
|
||||
variables: AlertingVariables;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import type { UploadProps } from "antd/es/upload";
|
|||
import React, { useState } from "react";
|
||||
import ProviderSpecificFields from "../add_model/provider_specific_fields";
|
||||
import { Providers, providerLogoMap } from "../provider_info_helpers";
|
||||
import { resetCredentialFormOnProviderChange } from "./credential_form_helpers";
|
||||
const { Link } = Typography;
|
||||
|
||||
interface AddCredentialsModalProps {
|
||||
|
|
@ -59,8 +60,7 @@ const AddCredentialsModal: React.FC<AddCredentialsModalProps> = ({ open, onCance
|
|||
<AntdSelect
|
||||
showSearch
|
||||
onChange={(value) => {
|
||||
setSelectedProvider(value as Providers);
|
||||
form.setFieldValue("custom_llm_provider", value);
|
||||
resetCredentialFormOnProviderChange(form, value as Providers, setSelectedProvider);
|
||||
}}
|
||||
>
|
||||
{Object.entries(Providers).map(([providerEnum, providerDisplayName]) => (
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import { useEffect, useState } from "react";
|
|||
import ProviderSpecificFields from "../add_model/provider_specific_fields";
|
||||
import { CredentialItem } from "../networking";
|
||||
import { Providers, providerLogoMap } from "../provider_info_helpers";
|
||||
import { resetCredentialFormOnProviderChange } from "./credential_form_helpers";
|
||||
const { Link } = Typography;
|
||||
|
||||
interface EditCredentialsModalProps {
|
||||
|
|
@ -92,8 +93,7 @@ export default function EditCredentialsModal({
|
|||
<AntdSelect
|
||||
showSearch
|
||||
onChange={(value) => {
|
||||
setSelectedProvider(value as Providers);
|
||||
form.setFieldValue("custom_llm_provider", value);
|
||||
resetCredentialFormOnProviderChange(form, value as Providers, setSelectedProvider);
|
||||
}}
|
||||
>
|
||||
{Object.entries(Providers).map(([providerEnum, providerDisplayName]) => (
|
||||
|
|
|
|||
|
|
@ -0,0 +1,84 @@
|
|||
import type { FormInstance } from "antd";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { Providers } from "../provider_info_helpers";
|
||||
import { resetCredentialFormOnProviderChange } from "./credential_form_helpers";
|
||||
|
||||
/**
|
||||
* Build a minimal FormInstance stub that records calls. We don't depend
|
||||
* on the full Antd API surface — only the three methods the helper uses.
|
||||
*/
|
||||
function makeFormStub(initialFields: Record<string, unknown> = {}) {
|
||||
const fields: Record<string, unknown> = { ...initialFields };
|
||||
const stub = {
|
||||
getFieldValue: vi.fn((key: string) => fields[key]),
|
||||
setFieldValue: vi.fn((key: string, value: unknown) => {
|
||||
fields[key] = value;
|
||||
}),
|
||||
resetFields: vi.fn(() => {
|
||||
Object.keys(fields).forEach((k) => delete fields[k]);
|
||||
}),
|
||||
};
|
||||
return { stub: stub as unknown as FormInstance, fields, calls: stub };
|
||||
}
|
||||
|
||||
describe("resetCredentialFormOnProviderChange", () => {
|
||||
it("clears all fields when switching providers", () => {
|
||||
// Simulate the OpenAI->Google AI Studio leak: api_base picked up
|
||||
// OpenAI's default value and the user typed a custom URL.
|
||||
const { stub, fields, calls } = makeFormStub({
|
||||
credential_name: "my-prod-key",
|
||||
custom_llm_provider: "OpenAI",
|
||||
api_base: "https://api.openai.com/v1",
|
||||
api_key: "sk-stale-openai-key",
|
||||
organization: "org-leak",
|
||||
});
|
||||
const setSelectedProvider = vi.fn();
|
||||
|
||||
resetCredentialFormOnProviderChange(stub, Providers.Google_AI_Studio, setSelectedProvider);
|
||||
|
||||
expect(calls.resetFields).toHaveBeenCalledTimes(1);
|
||||
// Provider-specific fields must be gone so the next render starts
|
||||
// from the new provider's default_value, not OpenAI's leftover.
|
||||
expect(fields.api_base).toBeUndefined();
|
||||
expect(fields.api_key).toBeUndefined();
|
||||
expect(fields.organization).toBeUndefined();
|
||||
});
|
||||
|
||||
it("preserves credential_name across the switch", () => {
|
||||
// credential_name is user-supplied metadata, not provider-specific.
|
||||
// The admin shouldn't have to retype it just because they re-picked
|
||||
// the provider.
|
||||
const { stub, fields } = makeFormStub({
|
||||
credential_name: "my-prod-key",
|
||||
custom_llm_provider: "OpenAI",
|
||||
api_base: "https://api.openai.com/v1",
|
||||
});
|
||||
|
||||
resetCredentialFormOnProviderChange(stub, Providers.Google_AI_Studio, vi.fn());
|
||||
|
||||
expect(fields.credential_name).toBe("my-prod-key");
|
||||
});
|
||||
|
||||
it("updates custom_llm_provider and selectedProvider state to the new value", () => {
|
||||
const { stub, fields } = makeFormStub({ credential_name: "x" });
|
||||
const setSelectedProvider = vi.fn();
|
||||
|
||||
resetCredentialFormOnProviderChange(stub, Providers.Google_AI_Studio, setSelectedProvider);
|
||||
|
||||
expect(fields.custom_llm_provider).toBe(Providers.Google_AI_Studio);
|
||||
expect(setSelectedProvider).toHaveBeenCalledExactlyOnceWith(Providers.Google_AI_Studio);
|
||||
});
|
||||
|
||||
it("does not call setFieldValue('credential_name', undefined) when the name was unset", () => {
|
||||
// Edge case: brand-new modal with no name typed yet. We shouldn't
|
||||
// explicitly write `undefined` back into the form (Antd treats that
|
||||
// as a touched empty field, triggering the "required" validation
|
||||
// prematurely).
|
||||
const { stub, calls } = makeFormStub({});
|
||||
|
||||
resetCredentialFormOnProviderChange(stub, Providers.Anthropic, vi.fn());
|
||||
|
||||
const credentialNameCalls = calls.setFieldValue.mock.calls.filter(([key]) => key === "credential_name");
|
||||
expect(credentialNameCalls).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
import type { FormInstance } from "antd";
|
||||
import { Providers } from "../provider_info_helpers";
|
||||
|
||||
/**
|
||||
* Reset the credential form when the user switches providers.
|
||||
*
|
||||
* Why: provider-specific fields (api_base, api_key, organization, ...)
|
||||
* share a single Antd Form state across providers. Without this reset,
|
||||
* the previous provider's values stick around — most visibly, OpenAI's
|
||||
* default `api_base` (https://api.openai.com/v1) carries over when the
|
||||
* user switches to Google AI Studio, overriding that provider's own
|
||||
* default_value.
|
||||
*
|
||||
* Strategy: blow away the whole form, then restore the provider-agnostic
|
||||
* fields (credential name + the new provider id) so the newly rendered
|
||||
* `ProviderSpecificFields` can apply its own defaults from a clean slate.
|
||||
*
|
||||
* The credential name is preserved because it's a user-supplied label
|
||||
* that shouldn't reset just because the admin re-selected a provider.
|
||||
*/
|
||||
export function resetCredentialFormOnProviderChange(
|
||||
form: FormInstance,
|
||||
newProvider: Providers,
|
||||
setSelectedProvider: (p: Providers) => void,
|
||||
): void {
|
||||
const preservedName = form.getFieldValue("credential_name");
|
||||
form.resetFields();
|
||||
if (preservedName !== undefined) {
|
||||
form.setFieldValue("credential_name", preservedName);
|
||||
}
|
||||
setSelectedProvider(newProvider);
|
||||
form.setFieldValue("custom_llm_provider", newProvider);
|
||||
}
|
||||
|
|
@ -285,6 +285,7 @@ export interface LiteLLMWellKnownUiConfig {
|
|||
auto_redirect_to_sso: boolean;
|
||||
admin_ui_disabled: boolean;
|
||||
sso_configured: boolean;
|
||||
hide_default_credentials_hint?: boolean;
|
||||
is_control_plane?: boolean;
|
||||
workers?: WorkerInfo[];
|
||||
}
|
||||
|
|
|
|||
29
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
29
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25016,6 +25016,10 @@ export interface components {
|
|||
cache_read_input_token_cost?: number | null;
|
||||
/** Cache Read Input Token Cost Above 200K Tokens */
|
||||
cache_read_input_token_cost_above_200k_tokens?: number | null;
|
||||
/** Cache Read Input Token Cost Above 200K Tokens Priority */
|
||||
cache_read_input_token_cost_above_200k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens Priority */
|
||||
cache_read_input_token_cost_above_272k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Flex */
|
||||
cache_read_input_token_cost_flex?: number | null;
|
||||
/** Cache Read Input Token Cost Priority */
|
||||
|
|
@ -25064,6 +25068,10 @@ export interface components {
|
|||
input_cost_per_token_above_128k_tokens?: number | null;
|
||||
/** Input Cost Per Token Above 200K Tokens */
|
||||
input_cost_per_token_above_200k_tokens?: number | null;
|
||||
/** Input Cost Per Token Above 200K Tokens Priority */
|
||||
input_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens Priority */
|
||||
input_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Batches */
|
||||
input_cost_per_token_batches?: number | null;
|
||||
/** Input Cost Per Token Cache Hit */
|
||||
|
|
@ -25137,6 +25145,10 @@ export interface components {
|
|||
output_cost_per_token_above_128k_tokens?: number | null;
|
||||
/** Output Cost Per Token Above 200K Tokens */
|
||||
output_cost_per_token_above_200k_tokens?: number | null;
|
||||
/** Output Cost Per Token Above 200K Tokens Priority */
|
||||
output_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens Priority */
|
||||
output_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Batches */
|
||||
output_cost_per_token_batches?: number | null;
|
||||
/** Output Cost Per Token Flex */
|
||||
|
|
@ -31108,6 +31120,11 @@ export interface components {
|
|||
admin_ui_disabled: boolean;
|
||||
/** Auto Redirect To Sso */
|
||||
auto_redirect_to_sso: boolean;
|
||||
/**
|
||||
* Hide Default Credentials Hint
|
||||
* @default false
|
||||
*/
|
||||
hide_default_credentials_hint: boolean;
|
||||
/**
|
||||
* Is Control Plane
|
||||
* @default false
|
||||
|
|
@ -32657,6 +32674,10 @@ export interface components {
|
|||
cache_read_input_token_cost?: number | null;
|
||||
/** Cache Read Input Token Cost Above 200K Tokens */
|
||||
cache_read_input_token_cost_above_200k_tokens?: number | null;
|
||||
/** Cache Read Input Token Cost Above 200K Tokens Priority */
|
||||
cache_read_input_token_cost_above_200k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Above 272K Tokens Priority */
|
||||
cache_read_input_token_cost_above_272k_tokens_priority?: number | null;
|
||||
/** Cache Read Input Token Cost Flex */
|
||||
cache_read_input_token_cost_flex?: number | null;
|
||||
/** Cache Read Input Token Cost Priority */
|
||||
|
|
@ -32705,6 +32726,10 @@ export interface components {
|
|||
input_cost_per_token_above_128k_tokens?: number | null;
|
||||
/** Input Cost Per Token Above 200K Tokens */
|
||||
input_cost_per_token_above_200k_tokens?: number | null;
|
||||
/** Input Cost Per Token Above 200K Tokens Priority */
|
||||
input_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Above 272K Tokens Priority */
|
||||
input_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Input Cost Per Token Batches */
|
||||
input_cost_per_token_batches?: number | null;
|
||||
/** Input Cost Per Token Cache Hit */
|
||||
|
|
@ -32778,6 +32803,10 @@ export interface components {
|
|||
output_cost_per_token_above_128k_tokens?: number | null;
|
||||
/** Output Cost Per Token Above 200K Tokens */
|
||||
output_cost_per_token_above_200k_tokens?: number | null;
|
||||
/** Output Cost Per Token Above 200K Tokens Priority */
|
||||
output_cost_per_token_above_200k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Above 272K Tokens Priority */
|
||||
output_cost_per_token_above_272k_tokens_priority?: number | null;
|
||||
/** Output Cost Per Token Batches */
|
||||
output_cost_per_token_batches?: number | null;
|
||||
/** Output Cost Per Token Flex */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue