chore(oss): litellm oss staging 120626 (#30292)

* feat(bedrock): add bedrock mantle gemma 4 models (#30264)

* feat(bedrock): add bedrock mantle gemma 4 models

* test(bedrock): harden mantle local cost fixture

* feat(responses): enable the responses API for the Tensormesh provider (#30209)

* feat(responses): enable the responses API for the Tensormesh provider

* Update litellm/llms/openai_like/providers.json

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(langfuse_otel): mark LLM spans as generations (#30250)

* fix(bedrock): stop stream_chunk_size leaking into invoke request bodies (#30240)

stream_chunk_size is a LiteLLM-internal knob for re-chunking the HTTP
response stream. The invoke transformations splat optional_params into the
provider request body without dropping it, and Bedrock rejects unknown
fields, so any bedrock/invoke request that sets the parameter fails with
ValidationException: stream_chunk_size: Extra inputs are not permitted.
Drop it in the invoke dispatcher (covers cohere, titan, mistral, meta,
ai21) and in the Claude messages-format request builder (the route used
for bedrock/invoke Anthropic models)

* fix(bedrock): stop buffering streamed tool-call argument deltas (#30231)

* fix(bedrock): stop buffering streamed tool-call argument deltas

Two issues made Bedrock tool-use streaming arrive as a single end-of-stream
burst through LiteLLM while plain text streamed fine.

First, the anthropic-beta allowlist mapped fine-grained-tool-streaming-2025-05-14
to null for bedrock and bedrock_converse, so the header was silently stripped.
Without that beta, Anthropic models on Bedrock buffer tool input server-side and
emit all toolUse.input deltas at once (verified against converse-stream and
invoke-with-response-stream directly). Bedrock accepts the beta via
additionalModelRequestFields.anthropic_beta, so it is now forwarded.

Second, the streaming reads re-chunked the AWS event stream with
iter_bytes(chunk_size=1024). httpx's ByteChunker only releases full 1024-byte
blocks, so the small early events (messageStart, contentBlockStart, first
deltas) sat in the buffer until enough bytes accumulated, pushing
time-to-first-byte from ~1.4s to ~8.5s on buffered tool-use streams. The
default is now no re-chunking; an explicit stream_chunk_size is still honored.

* test(bedrock): cover explicit stream_chunk_size on sync invoke path

* test(bedrock): cover stream_chunk_size plumbing through converse completion

* test(bedrock): cover stream_chunk_size default in legacy BedrockLLM streaming

* test(bedrock): merge converse handler tests into existing mapped test file

pytest imports test modules by basename in non-package test dirs, so the new
tests/test_litellm/llms/bedrock/chat/test_converse_handler.py collided with
the pre-existing tests/test_litellm/llms/chat/test_converse_handler.py and
broke collection in CI. Move the new tests into the existing file

* feat(otel): emit v2 cost breakdown + stamp tracer scope version (#30156)

Read the StandardLoggingPayload cost_breakdown into a typed LLMCost on
LLMCallSpanData and emit each component under litellm.cost.* (absent
components omitted, so spans stay sparse). Stamp litellm.__version__ as
the instrumentation scope version so every v2 span carries a
deterministic scope.version.

Tests under tests/test_litellm/integrations/otel/.

* fix(proxy): cancel in-flight upstream LLM request on client disconnect (opt-in) (#30223)

* fix(proxy): cancel in-flight upstream LLM request on client disconnect (opt-in)

On the non-streaming path, base_process_llm_request awaited the LLM call
with no disconnect monitoring; when the HTTP client went away the
upstream request kept running until completion or request_timeout (6000s
default), holding a backend slot (e.g. a vLLM GPU slot) for output
nobody would read

Add an opt-in general_settings.cancel_on_disconnect flag, default off,
so the default code path is unchanged. When enabled, a receive-based
watcher task observes http.disconnect and cancels the asyncio.gather
driving the upstream call. The resulting CancelledError is converted to
HTTPException 499 only when the disconnect event is set, so
server-initiated cancellations still propagate as-is. The 499 then flows
through _handle_llm_api_exception like any other failure, meaning
post_call_failure_hook still releases max_parallel_requests slots and
fires spend and alerting callbacks; it is logged at info level instead
of a full traceback

Also removes the dead check_request_disconnection helper in
proxy_server.py (zero call sites) along with its behavior-pin tests

Builds on the receive-based design from #25776

Addresses #13774. Re-fixes #22805 (regressed after the #14295 revert)

Co-authored-by: CreateRandom <18438707+CreateRandom@users.noreply.github.com>

* fix(proxy): scope 499 quiet logging to disconnects and harden watcher

Address the two P2 findings from the Greptile review on #30223. The
info-level logging in _log_llm_api_exception now applies only to the
disconnect-specific HTTPException (status 499 plus the shared
_CLIENT_DISCONNECT_DETAIL message), so any other 499 raised by hooks or
guardrails keeps its full traceback. The disconnect watcher now catches
exceptions from request.receive() (e.g. a transport reset) and logs a
warning instead of dying silently, making the degradation to no-op
visible; a test pins that the LLM call is not cancelled in that case

---------

Co-authored-by: kursad <kursad.lacin@brado.net>
Co-authored-by: CreateRandom <18438707+CreateRandom@users.noreply.github.com>

* fix(bedrock): grant aws-external-anthropic:* in OIDC session policy for claude_platform (#30200) (#30205)

The inline STS session policy passed to assume_role_with_web_identity
acts as an IAM PERMISSION CEILING — effective permissions are the
intersection of the role's identity policies and this policy. Any
action not listed is silently denied even when the IAM role grants it.

#27678 added the bedrock/claude_platform/<model> route but its
service-side action namespace is aws-external-anthropic:*, not
bedrock:*. Without a matching statement here, every claude_platform
request via OIDC (GCP federation, EKS Pod Identity webhook, etc.) 403s
with 'no session policy allows the aws-external-anthropic:CreateInference
action' — even with a fully permissive identity policy.

Add a second ClaudePlatformLiteLLM statement covering CreateInference,
CreateBatchInference, CancelBatchInference, DeleteBatchInference,
CountTokens, Get*, List*. Keep aws:SecureTransport=true parity with the
bedrock statement.

Static creds + IRSA flow through different code paths and are not
affected.

Fixes #30200

* fix(proxy): set Retry-After header on RouterRateLimitError 429 responses (#30098)

* Set Retry-After header on RouterRateLimitError responses

When all deployments for a model are in cooldown, the proxy returns a
429 whose cooldown timing is only available by parsing the error
message string. RouterRateLimitError already carries cooldown_time, so
expose it as a standard retry-after header in
_handle_llm_api_exception. The value is rounded up so clients never
retry before the cooldown window ends.

Fixes #27823.

* Set Retry-After after response-headers hook so cooldown wins

The cooldown-derived retry-after was assigned before the
post_call_response_headers_hook merge, so a callback returning a
retry-after key (including a stale or empty value) silently clobbered
it. Move the RouterRateLimitError block after the callback merge so the
cooldown value is authoritative for this error type.

* fix(router): route aspeech through async_function_with_fallbacks (#30104)

* fix(router): route aspeech through async_function_with_fallbacks

Router.aspeech selected a deployment and awaited litellm.aspeech
directly, so TTS requests got no retry on failure and no failover to
backup deployments; the except block only fired an exception alert and
re-raised. Every other router endpoint (acompletion, aembedding,
atranscription, arerank) already delegates to
async_function_with_fallbacks

Mirror the atranscription pattern: move deployment selection and the
litellm.aspeech call into a private _aspeech method, then have the
public aspeech set kwargs["original_function"] = self._aspeech and
await self.async_function_with_fallbacks(**kwargs). _aspeech also picks
up the shared _get_async_openai_model_client helper and the same
total/success/fail call accounting the sibling endpoints use

Fixes #27778.

* fix(router): apply deployment kwargs and rpm semaphore in _aspeech

Bring _aspeech fully in line with _atranscription: call
_update_kwargs_with_deployment so deployment metadata, model_info,
timeout, and default litellm params flow into the request, and wrap
the litellm.aspeech call with the max_parallel_requests semaphore plus
async_routing_strategy_pre_call_checks so TTS respects rpm limits the
same way the other router endpoints do

Also add a unit test that exercises _aspeech directly and asserts the
deployment metadata reaches the underlying call

* fix(slack_alerting): stop false-positive hanging request alerts for requests below the alerting threshold (#30106)

* fix(slack_alerting): skip hanging request alerts below the threshold

The hanging request check alerted on any cached request whose
completion status was not yet recorded, with no minimum age check.
Since the background loop runs every alerting_threshold / 2 seconds,
any request that happened to be in flight at a check fired a
"hanging - Ns+ request time" alert even if it was only seconds old,
producing a steady stream of false positives.

Add a created_at timestamp to HangingRequestData, stamped when the
request enters the hanging request cache, and skip requests younger
than alerting_threshold without evicting them, so a later check can
still alert if they never complete. Extend the cache TTL from
threshold + 60s to 1.5x threshold + 60s; with the age check, entries
only become alertable after threshold seconds, and the check period
is threshold / 2, so the old TTL could evict a genuinely hanging
request before any check saw it cross the threshold.

Fixes #27855.

* fix(slack_alerting): alert once per hanging request

The min-age gate stops false positives for young in-flight requests, but
a genuinely hanging request still re-alerted on every checker tick within
the cache TTL. With the wider TTL (1.5x threshold + 60s) that is 1-2 extra
Slack notifications per stuck request at the default 600s threshold.

Flag a HangingRequestData entry as alerted once its alert fires and skip
flagged entries on later ticks, so each hang produces exactly one alert.
The cache reference is mutated in place, so the TTL is untouched and still
handles cleanup. Adds a regression test asserting one alert across multiple
ticks.

Fixes #27855.

* fix(health): treat all-proxy-models keys as unrestricted in /health (#30087)

* fix(health): treat all-proxy-models keys as unrestricted in /health

A key granted all model permissions stores the literal
"all-proxy-models" marker in its models list. The /health access
filter compared that marker against real model_names, so the model
list filtered down to nothing and the WebUI health check returned
healthy_count=0, unhealthy_count=0 with HTTP 503. Skip the filter
(both the live path and the background-cache model_id scoping) when
the marker is present, matching how auth_checks treats
SpecialModelNames.all_proxy_models.

Fixes #29744.

* fix(health): resolve all-team-models sentinel to the team allowlist

Same failure shape as the all-proxy-models case: a key carrying the
literal "all-team-models" entry matches no real model_name, so the
/health access filter would zero out the model list. Resolve the
sentinel to the key's team models when team_id is set, matching
get_key_models in model_checks.py. Without a team_id the sentinel
stays unresolved and matches nothing, denying rather than widening
access, mirroring _resolve_key_models_for_auth_check.

* feat(proxy): auto-enable drop_params for Claude Code requests (#30218)

* feat(proxy): auto-enable drop_params for Claude Code requests

Claude Code identifies itself with a claude-cli/<version> user agent and
sends Anthropic-specific params (top_k, thinking, etc.) on every request.
When the proxy routes those requests to a non-Anthropic provider, the
unsupported params fail the call unless drop_params is configured. Detect
the Claude Code user agent in add_litellm_data_to_request and default
drop_params to true for those requests, without overriding an explicit
drop_params value sent by the caller.

* feat(proxy): respect operator litellm_settings drop_params over Claude Code default

An explicit drop_params in the operator's litellm_settings (true or false)
now suppresses the Claude Code user agent default, so an operator who
deliberately configured drop_params: false keeps strict param validation
for Claude Code clients too. The auto-default only fills the gap when
neither the request body nor the config sets a value.

* fix(snowflake): migrate to native endpoints with auto-routing for Claude models (#29964)

* fix(snowflake): migrate to native Cortex REST API endpoints

Replaces the legacy /api/v2/cortex/inference:complete endpoint with the
native OpenAI-compatible /api/v2/cortex/v1/chat/completions endpoint,
fixing error 390142 (Incoming request does not contain a valid payload)
when using model: snowflake/<model> in LiteLLM proxy.

Changes:
- litellm/llms/snowflake/chat/transformation.py: route to native
  /cortex/v1/chat/completions, remove Snowflake-specific tool_spec
  payload transformation, remove content_list response handling,
  add stream to supported params
- litellm/llms/snowflake/anthropic/transformation.py (new):
  SnowflakeCortexAnthropicConfig routes Claude models to /cortex/v1/messages
  with anthropic-version header and Anthropic->OpenAI response transform
- tests: 29 unit tests covering URL routing, auth headers, payload
  format, and response parsing

* fix(snowflake): map max_tokens to max_completion_tokens for native endpoint

* fix: handle multi-turn tool conversations and OpenAI→Anthropic tool format conversion

- _extract_system_and_messages now preserves tool_calls from assistant messages
  and converts them to Anthropic tool_use content blocks
- tool role messages are converted to user role with tool_result content blocks
  (as required by Anthropic Messages API)
- Added _transform_tools_to_anthropic() to convert OpenAI tool format
  (type/function/parameters) to Anthropic format (name/input_schema)
- Added comprehensive tests for multi-turn tool conversations

Addresses review feedback on PR #29964

* test: add coverage for malformed JSON and non-string tool arguments

* fix(tests): update chat transformation tests for native OpenAI-compatible endpoint

* style: apply black formatting

* fix: resolve mypy type errors in anthropic transformation

* fix: correct mypy type: ignore error codes (attr-defined)

* fix: use max_tokens instead of max_completion_tokens for Snowflake endpoint compatibility

* refactor: merge Anthropic config into unified SnowflakeConfig with auto-routing

- Remove separate SnowflakeCortexAnthropicConfig and anthropic/ directory
- SnowflakeConfig now auto-routes based on model name:
  - Claude models → /messages endpoint (Anthropic format)
  - All others → /chat/completions endpoint (OpenAI format)
- No new provider needed (stays as SNOWFLAKE = 'snowflake')
- Tool message transformation for Claude: tool_calls → tool_use blocks,
  tool role → user with tool_result
- OpenAI → Anthropic tool format conversion (parameters → input_schema)
- Addresses Greptile feedback about unwired SnowflakeCortexAnthropicConfig

* fix: use max_completion_tokens for /chat/completions (Snowflake deprecated max_tokens on this endpoint)

* fix(tests): update assertions for Claude auto-routing to /messages endpoint

* fix(snowflake): add tool_choice conversion and preserve max_completion_tokens in Anthropic path

* fix(snowflake): use ChatCompletionMessageToolCall objects and strip model prefix on OpenAI path

* fix(snowflake): collect multiple system messages to prevent guardrail override

* chore: remove committed .pyc files and add __pycache__ to .gitignore

* fix: remove unused Union import

* fix: restore original .gitignore (accidentally replaced in earlier commit)

* feat(snowflake): add streaming response handler for both Anthropic and OpenAI SSE formats

* fix: remove unused AsyncIterator and Iterator imports

* fix: add missing total_tokens to ChatCompletionUsageBlock

* fix(snowflake): coalesce consecutive tool results into single user message for Anthropic

* fix(snowflake): handle message_start event for streaming input_tokens tracking

* fix: evict last deleted model in multi-instance deployments (#28608)

* fix: evict last deleted model in multi-instance deployments

_delete_deployment had an early return when db_models was empty,
preventing eviction of the last deleted model during reconciliation.

- Remove len(db_models)==0 early return from _delete_deployment
- Return None (not []) from _get_models_from_db on DB failure so
  callers can distinguish a transient failure from a genuinely empty DB
- Guard _update_llm_router against None to skip updates on DB failure

Fixes #28443

* test: remove dead MagicMock assignment in type_mismatch test

* fix: update test to pass [] not None to _update_llm_router

test_ProxyConfig__update_llm_router_bad_proxy_logging_raises was passing
None as new_models to get through to the proxy_logging_obj check, but
the None guard we added now returns early before reaching that path.
Pass [] instead so the test exercises the intended AttributeError case.

Signed-off-by: Rudra Dudhat <contact.rdudhat@gmail.com>

* chore: regenerate API types to sync schema.d.ts with proxy OpenAPI spec

Signed-off-by: Rudra Dudhat <contact.rdudhat@gmail.com>

---------

Signed-off-by: Rudra Dudhat <contact.rdudhat@gmail.com>

* fix: invalidate Redis spend counter on /key/reset_spend (#29694)

* fix: set Redis spend counter to reset_to value on /key/reset_spend

Previously, the Redis spend counter was always set to 0.0 after a reset,
even when reset_to was a non-zero value (partial reset). This caused
the budget to be under-enforced for up to 60 seconds until the counter
expired and fell through to the DB.

Now the counter is set to the actual reset_to value, so partial resets
are reflected correctly and budget enforcement is consistent.

* test: update reset_key_spend test to match direct cache set

The implementation now sets spend_counter_cache directly instead of
calling _invalidate_spend_counter. Update the test to verify the
in_memory_cache.set_cache call with the correct key, value, and ttl.

---------

Co-authored-by: michaelxer <michaelxer@users.noreply.github.com>

* fix: add scaleway models pricing (#27659)

* fix: Add embeddings support for Scaleway provider

* fix: resolve merge conflicts

* fix(main): clarify backend route handling for Swagger static assets (#30196)

* fix(main): clarify backend route handling for Swagger static assets

* fix(allowlist): add BACKEND_MOUNT_PATHS for Swagger static assets

* fix(voyage): route multimodal embeddings to correct endpoint (#30193)

* fix(voyage): route multimodal embeddings to correct endpoint

* test(voyage): cover multimodal embedding edge cases

* test(voyage): cover api key fallback

* fix(voyage): raise early on missing api key and malformed image url

* test(voyage): cover utils routing and helper

* fix(voyage): route supported openai params for multimodal models

* style: apply black formatting

* fix(ui): infer Azure API version from API base (#30204)

* fix(ui): infer Azure API version from API base

* fix(ui): address Azure API version feedback

* Update litellm/llms/snowflake/chat/transformation.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* feat(datadog): add team-scoped Datadog callback support (#29947)

Enable teams to configure their own Datadog credentials via
POST /team/{team_id}/callback, following the same pattern as Langfuse.

* Merge pull request #29528 from aanchal22/litellm_byok-alias-merge

fix(proxy): atomic merge for team model aliases and team.models on BYOK create

* feat: add EmpirioLabs as an OpenAI-compatible provider (#30278)

Co-authored-by: Adam Dalloul <adam.d.developer@gmail.com>

* fix: resolve failing tests and lint in snowflake/team endpoints

- Black-format snowflake/chat/transformation.py to fix lint failure
- Update Anthropic config test to expect default max_tokens of 4096 (matches implementation)
- Add AsyncMock + execute_raw mock to team_model_add cache-refresh pin test
- Add model_dump mock and patch cache/logging in test_uses_atomic_array_append_with_dedup

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(test): update test_db_error_new_model_check for new _delete_deployment logic

_delete_deployment no longer short-circuits on empty db_models — it now
treats [] as a valid empty-DB state and proceeds to check config models.
Mock get_config to return the two router deployments so they appear in
combined_id_list and are protected, which matches the real-world scenario
where a DB error occurs but the models are config-backed.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* feat(proxy): register cancel_on_disconnect in ConfigGeneralSettings and config list (#30295)

* feat(proxy): register cancel_on_disconnect in ConfigGeneralSettings and config list

Follow-up to #30223 per maintainer review: documents the flag in
ConfigGeneralSettings with a short description and adds it to
allowed_args in get_config_list so the UI and /config/list expose it.
A test pins that /config/list returns the field with type Boolean,
which requires both registrations to be present

* chore(ui): regenerate schema.d.ts for cancel_on_disconnect

---------

Co-authored-by: kursad <kursad.lacin@brado.net>

* fix(datadog): never fall back to env DD_API_KEY for caller-supplied destinations

Team/key-scoped Datadog loggers could be pointed at an arbitrary dd_agent_host or
dd_site while omitting dd_api_key, causing the proxy's global DD_API_KEY to be sent
as the DD-API-KEY header to that destination. Gate the env-var fallback behind an
allow_env_credentials flag, set to False when the destination is caller-supplied,
mirroring the existing langfuse/langsmith pattern.

---------

Signed-off-by: Rudra Dudhat <contact.rdudhat@gmail.com>
Co-authored-by: Emerson Gomes <emerson.gomes@thalesgroup.com>
Co-authored-by: daitran-tensormesh <dai@tensormesh.ai>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Muspi Merol <me@promplate.dev>
Co-authored-by: fangkang <fangkangm@gmail.com>
Co-authored-by: Chris Hoogeboom <chris.hoogeboom@gmail.com>
Co-authored-by: kursadlacin <kursadlacin@gmail.com>
Co-authored-by: kursad <kursad.lacin@brado.net>
Co-authored-by: CreateRandom <18438707+CreateRandom@users.noreply.github.com>
Co-authored-by: hcl <chenglunhu@gmail.com>
Co-authored-by: Filippo Menghi <113345637+Cyberfilo@users.noreply.github.com>
Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: sfc-gh-nashukla <navnit.shukla@snowflake.com>
Co-authored-by: Rudra Dudhat <contact.rdudhat@gmail.com>
Co-authored-by: Michael <52305679+michaelxer@users.noreply.github.com>
Co-authored-by: michaelxer <michaelxer@users.noreply.github.com>
Co-authored-by: Quentin Champenois <26109239+Quentinchampenois@users.noreply.github.com>
Co-authored-by: mauriceberentsen <mauriceberentsen@live.nl>
Co-authored-by: lost9999 <56498264+lost9999@users.noreply.github.com>
Co-authored-by: GaetanVDB07 <86427581+GaetanVDB07@users.noreply.github.com>
Co-authored-by: Aanchal Khandelwal <aan2210khandelwal@gmail.com>
Co-authored-by: Adam Dalloul <adam_dalloul@icloud.com>
Co-authored-by: Adam Dalloul <adam.d.developer@gmail.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Sameer Kankute 2026-06-12 22:19:25 +05:30 committed by GitHub
parent 7d1f68e72a
commit 079c136742
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
69 changed files with 5055 additions and 633 deletions

View file

@ -20,7 +20,11 @@ DatabaseURLSettings.from_env().apply_to_env()
from litellm.proxy.proxy_server import app
from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES
from backend.routes.allowlist import (
BACKEND_EXACT_PATHS,
BACKEND_MOUNT_PATHS,
BACKEND_PATH_PREFIXES,
)
def _is_backend_route(route) -> bool:
@ -29,8 +33,9 @@ def _is_backend_route(route) -> bool:
if path is None:
return False
if isinstance(route, Mount):
# Static UI mounts are served by the dedicated UI container, not here.
return False
# The dashboard UI static mounts are served by the dedicated UI container.
# Only Mounts in the backend allowlist (e.g. swagger docs) remain on backend.
return path in BACKEND_MOUNT_PATHS
if path in BACKEND_EXACT_PATHS:
return True
return any(path.startswith(prefix) for prefix in BACKEND_PATH_PREFIXES)

View file

@ -133,3 +133,9 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
"/fallback/login",
}
)
BACKEND_MOUNT_PATHS: frozenset[str] = frozenset(
{
"/swagger", # API documentation static assets belong to the backend
}
)

View file

@ -1731,6 +1731,9 @@ if TYPE_CHECKING:
from .llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig,
)
from .llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig as VoyageMultimodalEmbeddingConfig,
)
from .llms.infinity.embedding.transformation import (
InfinityEmbeddingConfig as InfinityEmbeddingConfig,
)

View file

@ -223,6 +223,7 @@ LLM_CONFIG_NAMES = (
"GenAIHubOrchestrationConfig",
"VoyageEmbeddingConfig",
"VoyageContextualEmbeddingConfig",
"VoyageMultimodalEmbeddingConfig",
"InfinityEmbeddingConfig",
"PerplexityEmbeddingConfig",
"AzureAIStudioConfig",
@ -903,6 +904,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.voyage.embedding.transformation_contextual",
"VoyageContextualEmbeddingConfig",
),
"VoyageMultimodalEmbeddingConfig": (
".llms.voyage.embedding.transformation_multimodal",
"VoyageMultimodalEmbeddingConfig",
),
"InfinityEmbeddingConfig": (
".llms.infinity.embedding.transformation",
"InfinityEmbeddingConfig",

View file

@ -75,7 +75,7 @@
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"fine-grained-tool-streaming-2025-05-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": null,
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,
@ -106,7 +106,7 @@
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"fine-grained-tool-streaming-2025-05-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": null,
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,

View file

@ -8,6 +8,7 @@ Notes:
"""
import asyncio
import time
from typing import TYPE_CHECKING, Any, Optional
import litellm
@ -36,11 +37,15 @@ class AlertingHangingRequestCheck:
slack_alerting_object: SlackAlerting,
):
self.slack_alerting_object = slack_alerting_object
# checks run every alerting_threshold / 2 seconds, so entries must
# stay cached for at least 1.5x the threshold to guarantee a check
# happens after they cross it
self.hanging_request_cache_ttl = int(
self.slack_alerting_object.alerting_threshold * 1.5
+ HANGING_ALERT_BUFFER_TIME_SECONDS
)
self.hanging_request_cache = InMemoryCache(
default_ttl=int(
self.slack_alerting_object.alerting_threshold
+ HANGING_ALERT_BUFFER_TIME_SECONDS
),
default_ttl=self.hanging_request_cache_ttl,
)
async def add_request_to_hanging_request_check(
@ -76,10 +81,7 @@ class AlertingHangingRequestCheck:
await self.hanging_request_cache.async_set_cache(
key=hanging_request_data.request_id,
value=hanging_request_data,
ttl=int(
self.slack_alerting_object.alerting_threshold
+ HANGING_ALERT_BUFFER_TIME_SECONDS
),
ttl=self.hanging_request_cache_ttl,
)
return
@ -111,6 +113,9 @@ class AlertingHangingRequestCheck:
if hanging_request_data is None:
continue
if hanging_request_data.alerted:
continue
request_status = (
await proxy_logging_obj.internal_usage_cache.async_get_cache(
key="request_status:{}".format(hanging_request_data.request_id),
@ -127,12 +132,21 @@ class AlertingHangingRequestCheck:
)
continue
request_age_seconds = time.time() - hanging_request_data.created_at
if request_age_seconds < self.slack_alerting_object.alerting_threshold:
# in-flight but below the alerting threshold; keep it cached
# so a later check can alert if it never completes
continue
################
# Send the Alert on Slack
################
await self.send_hanging_request_alert(
hanging_request_data=hanging_request_data
)
# flag so the entry is skipped on later ticks; one alert per hang,
# with the existing TTL still handling cleanup
hanging_request_data.alerted = True
return

View file

@ -92,12 +92,26 @@ class DataDogLogger(
# Class variables or attributes
def __init__(
self,
dd_api_key: Optional[str] = None,
dd_site: Optional[str] = None,
dd_agent_host: Optional[str] = None,
dd_agent_port: Optional[str] = None,
allow_env_credentials: bool = True,
**kwargs,
):
"""
Initializes the datadog logger, checks if the correct env variables are set
Required environment variables (Direct API):
Args:
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True.
dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var.
dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var.
dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var.
allow_env_credentials: When False, the API key is never read from DD_API_KEY env var. Set to
False for team/key-scoped loggers whose destination (dd_agent_host/dd_site) is caller-supplied,
so the proxy's global DD_API_KEY is never sent to an untrusted host.
Required environment variables (Direct API) when kwargs not provided:
`DD_API_KEY` - your datadog api key
`DD_SITE` - your datadog site, example = `"us5.datadoghq.com"`
@ -130,12 +144,21 @@ class DataDogLogger(
)
# Configure DataDog endpoint (Agent or Direct API)
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")
if dd_agent_host:
self._configure_dd_agent(dd_agent_host=dd_agent_host)
# Prefer explicit kwargs, then fall back to env vars
resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST")
if resolved_agent_host:
self._configure_dd_agent(
dd_agent_host=resolved_agent_host,
dd_agent_port=dd_agent_port,
dd_api_key=dd_api_key,
allow_env_credentials=allow_env_credentials,
)
else:
self._configure_dd_direct_api()
self._configure_dd_direct_api(
dd_api_key=dd_api_key,
dd_site=dd_site,
allow_env_credentials=allow_env_credentials,
)
# Optional override for testing
dd_base_url = get_datadog_base_url_from_env()
@ -172,34 +195,60 @@ class DataDogLogger(
).model_dump()
return dict_datadog_params
def _configure_dd_agent(self, dd_agent_host: str) -> None:
def _configure_dd_agent(
self,
dd_agent_host: str,
dd_agent_port: Optional[str] = None,
dd_api_key: Optional[str] = None,
allow_env_credentials: bool = True,
) -> None:
"""
Configure DataDog Agent for log forwarding
Args:
dd_agent_host: Hostname or IP of DataDog agent
dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518).
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True. Optional when using agent.
allow_env_credentials: When False, never read the API key from DD_API_KEY env var.
"""
dd_agent_port = os.getenv(
resolved_port = dd_agent_port or os.getenv(
"LITELLM_DD_AGENT_PORT", "10518"
) # default port for logs
self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs"
self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent
self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs"
self.DD_API_KEY = dd_api_key or (
os.getenv("DD_API_KEY") if allow_env_credentials else None
) # Optional when using agent
verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}")
def _configure_dd_direct_api(self) -> None:
def _configure_dd_direct_api(
self,
dd_api_key: Optional[str] = None,
dd_site: Optional[str] = None,
allow_env_credentials: bool = True,
) -> None:
"""
Configure direct DataDog API connection
Args:
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True.
dd_site: Datadog site. Falls back to DD_SITE env var.
allow_env_credentials: When False, never read the API key from DD_API_KEY env var.
Raises:
Exception: If required environment variables are not set
Exception: If required credentials are not provided via args or env vars
"""
if os.getenv("DD_API_KEY", None) is None:
resolved_api_key = dd_api_key or (
os.getenv("DD_API_KEY") if allow_env_credentials else None
)
resolved_site = dd_site or os.getenv("DD_SITE")
if resolved_api_key is None:
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>")
if os.getenv("DD_SITE", None) is None:
if resolved_site is None:
raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>")
self.DD_API_KEY = os.getenv("DD_API_KEY")
self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs"
self.DD_API_KEY = resolved_api_key
self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs"
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""

View file

@ -0,0 +1,124 @@
"""
DataDog Team Handler
Used to get the DataDogLogger for a given request.
Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler.
"""
from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
from .datadog import DataDogLogger
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
else:
DynamicLoggingCache = Any
class DatadogLoggingConfig(TypedDict):
dd_api_key: Optional[str]
dd_site: Optional[str]
dd_agent_host: Optional[str]
dd_agent_port: Optional[str]
class DataDogHandler:
@staticmethod
def get_datadog_logger_for_request(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
in_memory_dynamic_logger_cache: DynamicLoggingCache,
) -> DataDogLogger:
"""
Get a team-scoped DataDogLogger for a given request.
Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache,
keyed by the team's DD credentials. Each unique set of credentials gets its own
logger instance with its own batch/flush loop.
Note: This handler is only called when team-scoped DD credentials are present.
The global (env-var based) DataDogLogger is managed separately by
_init_custom_logger_compatible_class via _in_memory_loggers.
"""
_credentials = DataDogHandler.get_dynamic_datadog_logging_config(
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
credentials_dict = dict(_credentials)
# check if datadog logger is already cached
temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache(
credentials=credentials_dict, service_name="datadog"
)
# if not cached, create a new datadog logger and cache it
if temp_datadog_logger is None:
temp_datadog_logger = (
DataDogHandler._create_datadog_logger_from_credentials(
credentials=credentials_dict,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
)
return temp_datadog_logger
@staticmethod
def _create_datadog_logger_from_credentials(
credentials: Dict,
in_memory_dynamic_logger_cache: DynamicLoggingCache,
) -> DataDogLogger:
"""
Create a DataDogLogger from the credentials and cache it.
"""
# When the destination is caller-supplied (dd_agent_host/dd_site), never fall back to the
# proxy's DD_API_KEY env var, otherwise it would be sent to a team-controlled host.
allow_env_credentials = (
credentials.get("dd_agent_host") is None
and credentials.get("dd_site") is None
)
datadog_logger = DataDogLogger(
dd_api_key=credentials.get("dd_api_key"),
dd_site=credentials.get("dd_site"),
dd_agent_host=credentials.get("dd_agent_host"),
dd_agent_port=credentials.get("dd_agent_port"),
allow_env_credentials=allow_env_credentials,
)
in_memory_dynamic_logger_cache.set_cache(
credentials=credentials,
service_name="datadog",
logging_obj=datadog_logger,
)
verbose_logger.debug(
"Datadog: Created and cached new DataDogLogger for team-scoped credentials"
)
return datadog_logger
@staticmethod
def get_dynamic_datadog_logging_config(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> DatadogLoggingConfig:
"""
Get the Datadog logging config for a given request from dynamic params.
"""
return DatadogLoggingConfig(
dd_api_key=standard_callback_dynamic_params.get("dd_api_key"),
dd_site=standard_callback_dynamic_params.get("dd_site"),
dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"),
dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"),
)
@staticmethod
def _dynamic_datadog_credentials_are_passed(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> bool:
"""
Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params.
"""
if (
standard_callback_dynamic_params.get("dd_api_key") is not None
or standard_callback_dynamic_params.get("dd_site") is not None
or standard_callback_dynamic_params.get("dd_agent_host") is not None
):
return True
return False

View file

@ -43,6 +43,7 @@ class LangfuseOtelLogger(OpenTelemetry):
"""
_utils.set_attributes(span, kwargs, response_obj, LangfuseLLMObsOTELAttributes)
span.set_attribute("langfuse.observation.type", "generation")
#########################################################
# Set Langfuse specific attributes

View file

@ -63,6 +63,20 @@ class GenAIMapper:
# routing) onto the boundary-born LLM span — stamp it directly here.
LiteLLM.PROVIDER_MODEL: lambda d: d.identity.provider_model or None,
f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost,
# Per-component cost breakdown (from the StandardLoggingPayload
# ``cost_breakdown``). Each component is omitted when the source didn't
# report it, so spans stay sparse rather than carrying zeros.
f"{LiteLLM.COST_PREFIX}input": lambda d: d.cost.input,
f"{LiteLLM.COST_PREFIX}output": lambda d: d.cost.output,
f"{LiteLLM.COST_PREFIX}cache_read": lambda d: d.cost.cache_read,
f"{LiteLLM.COST_PREFIX}cache_creation": lambda d: d.cost.cache_creation,
f"{LiteLLM.COST_PREFIX}tool_usage": lambda d: d.cost.tool_usage,
f"{LiteLLM.COST_PREFIX}original": lambda d: d.cost.original,
f"{LiteLLM.COST_PREFIX}discount_amount": lambda d: d.cost.discount_amount,
f"{LiteLLM.COST_PREFIX}discount_percent": lambda d: d.cost.discount_percent,
f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount,
f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent,
f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount,
LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming,
}

View file

@ -34,6 +34,7 @@ __all__ = [
"RequestIdentity",
"GuardrailSpanData",
"LLMCallSpanData",
"LLMCost",
"LLMRequestParams",
"LLMUsage",
"MCPToolCallSpanData",
@ -91,6 +92,49 @@ class LLMUsage:
total_tokens: int | None = None
@dataclass(frozen=True)
class LLMCost:
"""Per-component cost breakdown, from the StandardLoggingPayload
``cost_breakdown`` (``litellm.types.utils.CostBreakdown``).
Each field is the USD cost of one component, or ``None`` when the source did
not report it so the mapper omits absent components instead of emitting 0.
The final (post-discount/post-margin) total is carried separately on
``LLMCallSpanData.response_cost``. Free-form ``additional_costs`` are not
surfaced here: span attributes are scalar and there is no agreed key shape
for them yet.
"""
input: float | None = None
output: float | None = None
cache_read: float | None = None
cache_creation: float | None = None
tool_usage: float | None = None
original: float | None = None
discount_amount: float | None = None
discount_percent: float | None = None
margin_fixed_amount: float | None = None
margin_percent: float | None = None
margin_total_amount: float | None = None
@classmethod
def from_breakdown(cls, breakdown: Mapping[str, object] | None) -> "LLMCost":
b = breakdown or {}
return cls(
input=as_float(b.get("input_cost")),
output=as_float(b.get("output_cost")),
cache_read=as_float(b.get("cache_read_cost")),
cache_creation=as_float(b.get("cache_creation_cost")),
tool_usage=as_float(b.get("tool_usage_cost")),
original=as_float(b.get("original_cost")),
discount_amount=as_float(b.get("discount_amount")),
discount_percent=as_float(b.get("discount_percent")),
margin_fixed_amount=as_float(b.get("margin_fixed_amount")),
margin_percent=as_float(b.get("margin_percent")),
margin_total_amount=as_float(b.get("margin_total_amount")),
)
@dataclass(frozen=True)
class SpanError:
error_type: str | None = None
@ -255,6 +299,7 @@ class LLMCallSpanData:
server: ServerInfo | None
identity: RequestIdentity
is_streaming: bool | None = None
cost: LLMCost = field(default_factory=LLMCost)
tools: tuple[ToolDefinition, ...] = ()
# Raw messages and response, needed by vendor mappers (OpenInference,
# Langfuse, Weave) that stamp message-level attributes. ``messages_in`` is
@ -302,6 +347,9 @@ class LLMCallSpanData:
finish_reasons=finish_reasons,
error=_parse_error(payload),
response_cost=as_float(payload.get("response_cost")),
cost=LLMCost.from_breakdown(
cast("Mapping[str, object] | None", payload.get("cost_breakdown"))
),
server=ServerInfo.from_api_base(context.api_base),
identity=context.identity,
is_streaming=as_bool(payload.get("stream")),

View file

@ -17,6 +17,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
)
from opentelemetry.trace import Span, SpanKind, Tracer
from litellm._version import version as litellm_version
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.integrations.otel.model.semconv import LiteLLM
from litellm.integrations.otel.model.spans import LiteLLMSpanKind
@ -207,7 +208,10 @@ def build_tracer_provider(
def get_tracer(provider: TracerProvider, name: str = "litellm") -> Tracer:
return provider.get_tracer(name)
# Stamp the instrumentation scope with the LiteLLM package version so every
# emitted span carries a deterministic ``scope.version`` (the standard OTel
# location for the emitting library's version) for downstream consumers.
return provider.get_tracer(name, litellm_version)
def in_memory_provider(

View file

@ -295,6 +295,15 @@ def get_supported_openai_params( # noqa: PLR0915
elif custom_llm_provider == "predibase":
return litellm.PredibaseConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "voyage":
if (
request_type == "embeddings"
and litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model)
):
return (
litellm.VoyageMultimodalEmbeddingConfig().get_supported_openai_params(
model=model
)
)
return litellm.VoyageEmbeddingConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "infinity":
return litellm.InfinityEmbeddingConfig().get_supported_openai_params(

View file

@ -53,11 +53,19 @@ _supported_callback_params = [
"braintrust_host",
"slack_webhook_url",
"lunary_public_key",
"dd_api_key",
"dd_site",
"dd_agent_host",
"dd_agent_port",
]
_request_blocked_callback_params = {
"gcs_bucket_name",
"gcs_path_service_account",
"dd_api_key",
"dd_site",
"dd_agent_host",
"dd_agent_port",
}

View file

@ -381,13 +381,14 @@ class Logging(LiteLLMLoggingBaseClass):
List[Union[str, Callable, CustomLogger]]
] = dynamic_async_failure_callbacks
# Process dynamic callbacks
self.process_dynamic_callbacks()
## DYNAMIC LANGFUSE / GCS / logging callback KEYS ##
self.standard_callback_dynamic_params: StandardCallbackDynamicParams = (
self.initialize_standard_callback_dynamic_params(kwargs)
)
# Process dynamic callbacks (after standard_callback_dynamic_params is initialized,
# so team-scoped credentials are available for callback initialization)
self.process_dynamic_callbacks()
self.standard_built_in_tools_params: StandardBuiltInToolsParams = (
self.initialize_standard_built_in_tools_params(kwargs)
)
@ -482,8 +483,21 @@ class Logging(LiteLLMLoggingBaseClass):
isinstance(callback, str)
and callback in litellm._known_custom_logger_compatible_callbacks
):
# For callbacks that support team-scoped credentials (e.g. datadog),
# pass only the relevant dynamic params as custom_logger_init_args.
_custom_logger_init_args: Optional[dict] = None
if callback == "datadog":
_custom_logger_init_args = {
k: v
for k, v in self.standard_callback_dynamic_params.items()
if k.startswith("dd_")
}
callback_class = _init_custom_logger_compatible_class(
callback, internal_usage_cache=None, llm_router=None # type: ignore
callback, # type: ignore[arg-type]
internal_usage_cache=None,
llm_router=None, # type: ignore
custom_logger_init_args=_custom_logger_init_args,
)
if callback_class is not None:
processed_list.append(callback_class)
@ -3941,6 +3955,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_prometheus_logger)
return _prometheus_logger # type: ignore
elif logging_integration == "datadog":
# Check if team-scoped credentials are provided
_dd_api_key = custom_logger_init_args.get("dd_api_key")
_dd_site = custom_logger_init_args.get("dd_site")
_dd_agent_host = custom_logger_init_args.get("dd_agent_host")
_dd_agent_port = custom_logger_init_args.get("dd_agent_port")
if _dd_api_key or _dd_site or _dd_agent_host:
# Team-scoped credentials: use DynamicLoggingCache for per-credential isolation
from litellm.integrations.datadog.datadog_team_handler import (
DataDogHandler,
)
return DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=custom_logger_init_args, # type: ignore
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
# Global (env-var based): reuse cached instance
for callback in _in_memory_loggers:
if isinstance(callback, DataDogLogger):
return callback # type: ignore

View file

@ -861,14 +861,58 @@ class BaseAWSLLM:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client("sts", **sts_client_kwargs)
# The session policy is an IAM PERMISSION CEILING — effective
# permissions are the intersection of the role's identity policies
# and this policy. Any action not listed here is silently denied
# even when the IAM role grants it. So every Bedrock route we
# support needs a matching action statement, or it 403s on OIDC
# auth only (static creds + IRSA take other code paths).
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
bedrock_session_policy = {
"Version": "2012-10-17",
"Statement": [
{
"Sid": "BedrockLiteLLM",
"Effect": "Allow",
"Action": [
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
"bedrock:ApplyGuardrail",
"bedrock:GetGuardrail",
"bedrock:ListGuardrails",
],
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
# Claude Platform on AWS (added by #27678 for the
# ``bedrock/claude_platform/<model>`` route) lives under
# a separate IAM action namespace; without these entries
# the OIDC path 403s on every claude_platform request
# even with a fully permissive identity policy (#30200).
{
"Sid": "ClaudePlatformLiteLLM",
"Effect": "Allow",
"Action": [
"aws-external-anthropic:CreateInference",
"aws-external-anthropic:CreateBatchInference",
"aws-external-anthropic:CancelBatchInference",
"aws-external-anthropic:DeleteBatchInference",
"aws-external-anthropic:CountTokens",
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
],
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
],
}
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name,
"WebIdentityToken": oidc_token,
"DurationSeconds": 3600,
"Policy": '{"Version":"2012-10-17","Statement":[{"Sid":"BedrockLiteLLM","Effect":"Allow","Action":["bedrock:InvokeModel","bedrock:InvokeModelWithResponseStream","bedrock:ApplyGuardrail","bedrock:GetGuardrail","bedrock:ListGuardrails"],"Resource":"*","Condition":{"Bool":{"aws:SecureTransport":"true"}}}]}',
"Policy": json.dumps(bedrock_session_policy, separators=(",", ":")),
}
# Add ExternalId parameter if provided

View file

@ -32,7 +32,7 @@ def make_sync_call(
logging_obj: LiteLLMLoggingObject,
json_mode: Optional[bool] = False,
fake_stream: bool = False,
stream_chunk_size: int = 1024,
stream_chunk_size: Optional[int] = None,
):
if client is None:
client = _get_httpx_client() # Create a new client if none provided
@ -108,7 +108,7 @@ class BedrockConverseLLM(BaseAWSLLM):
fake_stream: bool = False,
json_mode: Optional[bool] = False,
api_key: Optional[str] = None,
stream_chunk_size: int = 1024,
stream_chunk_size: Optional[int] = None,
) -> CustomStreamWrapper:
request_data = await litellm.AmazonConverseConfig()._async_transform_request(
model=model,
@ -268,7 +268,7 @@ class BedrockConverseLLM(BaseAWSLLM):
):
## SETUP ##
stream = optional_params.pop("stream", None)
stream_chunk_size = optional_params.pop("stream_chunk_size", 1024)
stream_chunk_size = optional_params.pop("stream_chunk_size", None)
unencoded_model_id = optional_params.pop("model_id", None)
fake_stream = optional_params.pop("fake_stream", False)
json_mode = optional_params.get("json_mode", False)

View file

@ -197,7 +197,7 @@ async def make_call(
fake_stream: bool = False,
json_mode: Optional[bool] = False,
bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None,
stream_chunk_size: int = 1024,
stream_chunk_size: Optional[int] = None,
):
try:
if client is None:
@ -294,7 +294,7 @@ def make_sync_call(
fake_stream: bool = False,
json_mode: Optional[bool] = False,
bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None,
stream_chunk_size: int = 1024,
stream_chunk_size: Optional[int] = None,
):
try:
if client is None:
@ -790,7 +790,7 @@ class BedrockLLM(BaseAWSLLM):
## SETUP ##
stream = optional_params.pop("stream", None)
stream_chunk_size = optional_params.pop("stream_chunk_size", 1024)
stream_chunk_size = optional_params.pop("stream_chunk_size", None)
provider = self.get_bedrock_invoke_provider(model)
modelId = self.get_bedrock_model_id(
@ -1203,7 +1203,7 @@ class BedrockLLM(BaseAWSLLM):
extra_headers: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[AsyncHTTPHandler] = None,
stream_chunk_size: int = 1024,
stream_chunk_size: Optional[int] = None,
) -> Union[ModelResponse, CustomStreamWrapper]:
transformed_request = (
await litellm.AmazonAnthropicClaudeConfig().async_transform_request(
@ -1350,7 +1350,7 @@ class BedrockLLM(BaseAWSLLM):
logger_fn=None,
headers={},
client: Optional[AsyncHTTPHandler] = None,
stream_chunk_size: int = 1024,
stream_chunk_size: Optional[int] = None,
) -> CustomStreamWrapper:
# The call is not made here; instead, we prepare the necessary objects for the stream.

View file

@ -215,6 +215,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
anthropic_request.pop("model", None)
anthropic_request.pop("stream", None)
anthropic_request.pop("stream_chunk_size", None)
output_format = anthropic_request.pop("output_format", None)
output_config_format = pop_bedrock_invoke_output_config_format(
anthropic_request

View file

@ -150,6 +150,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
) -> dict:
## SETUP ##
stream = optional_params.pop("stream", None)
optional_params.pop("stream_chunk_size", None)
custom_prompt_dict: dict = litellm_params.pop("custom_prompt_dict", None) or {}
hf_model_name = litellm_params.get("hf_model_name", None)

View file

@ -131,7 +131,8 @@
"base_class": "openai_gpt",
"param_mappings": {
"max_completion_tokens": "max_tokens"
}
},
"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]
},
"parasail": {
"base_url": "https://api.parasail.io/v1",
@ -141,5 +142,14 @@
"special_handling": {
"force_store_false": true
}
},
"empiriolabs": {
"base_url": "https://api.empiriolabs.ai/v1",
"api_key_env": "EMPIRIOLABS_API_KEY",
"api_base_env": "EMPIRIOLABS_API_BASE",
"param_mappings": {
"max_completion_tokens": "max_tokens"
},
"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]
}
}

View file

@ -1,17 +1,32 @@
"""
Support for Snowflake REST API
Snowflake Cortex REST API Chat Transformation
Routes to native Cortex REST API endpoints based on model:
- Claude models POST /api/v2/cortex/v1/messages (Anthropic format)
- All other models POST /api/v2/cortex/v1/chat/completions (OpenAI format)
Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
"""
import json
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import httpx
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.types.utils import (
ChatCompletionMessageToolCall,
ChatCompletionUsageBlock,
Choices,
Function,
GenericStreamingChunk,
Message,
ModelResponse,
Usage,
)
from ...base_llm.base_model_iterator import BaseModelResponseIterator
from ...openai_like.chat.transformation import OpenAIGPTConfig
from ..utils import SnowflakeBaseConfig
if TYPE_CHECKING:
@ -21,69 +36,343 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = Any
ANTHROPIC_VERSION = "2023-06-01"
_CLAUDE_MODEL_PREFIXES = (
"claude-",
"claude_",
)
def _is_claude_model(model: str) -> bool:
"""Return True if model name (after stripping snowflake/ prefix) is a Claude model."""
name = model.lower().removeprefix("snowflake/")
return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES)
class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
"""
Reference: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api
Snowflake Cortex REST API unified provider.
Snowflake Cortex LLM REST API supports function calling with specific models (e.g., Claude 3.5 Sonnet).
This config handles transformation between OpenAI format and Snowflake's tool_spec format.
Auto-routes based on model name:
- Claude models /api/v2/cortex/v1/messages (Anthropic Messages format)
- All others /api/v2/cortex/v1/chat/completions (OpenAI format)
Auth:
PAT: api_key="pat/<token>" X-Snowflake-Authorization-Token-Type: PROGRAMMATIC_ACCESS_TOKEN
JWT: api_key="<jwt>" X-Snowflake-Authorization-Token-Type: KEYPAIR_JWT
"""
@classmethod
def get_config(cls):
return super().get_config()
def _transform_tool_calls_from_snowflake_to_openai(
self, content_list: List[Dict[str, Any]]
) -> Tuple[str, Optional[List[ChatCompletionMessageToolCall]]]:
def get_supported_openai_params(self, model: str) -> List[str]:
params = [
"temperature",
"max_tokens",
"max_completion_tokens",
"top_p",
"stream",
"tools",
"tool_choice",
]
if _is_claude_model(model):
params.append("thinking")
return params
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
api_base = self._get_api_base(api_base, optional_params)
if _is_claude_model(model):
return f"{api_base}/cortex/v1/messages"
return f"{api_base}/cortex/v1/chat/completions"
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
headers = super().validate_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
api_base=api_base,
)
if _is_claude_model(model):
headers["anthropic-version"] = ANTHROPIC_VERSION
return headers
def _transform_tools_to_anthropic(self, tools: List[Dict]) -> List[Dict]:
"""
Transform Snowflake tool calls to OpenAI format.
Convert tools from OpenAI format to Anthropic format.
Args:
content_list: Snowflake's content_list array containing text and tool_use items
OpenAI: {"type": "function", "function": {"name": ..., "parameters": {...}}}
Anthropic: {"name": ..., "description": ..., "input_schema": {...}}
"""
anthropic_tools = []
for tool in tools:
if tool.get("type") == "function" and "function" in tool:
func = tool["function"]
anthropic_tool: Dict[str, Any] = {
"name": func.get("name", ""),
}
if "description" in func:
anthropic_tool["description"] = func["description"]
if "parameters" in func:
anthropic_tool["input_schema"] = func["parameters"]
else:
anthropic_tool["input_schema"] = {
"type": "object",
"properties": {},
}
anthropic_tools.append(anthropic_tool)
else:
anthropic_tools.append(tool)
return anthropic_tools
Returns:
Tuple of (text_content, tool_calls)
def _extract_system_and_messages(
self, messages: List[AllMessageValues]
) -> tuple[Optional[str], List[Dict]]:
"""
Split messages into system prompt and conversation turns for Anthropic format.
Snowflake format in content_list:
{
"type": "tool_use",
"tool_use": {
"tool_use_id": "tooluse_...",
"name": "get_weather",
"input": {"location": "Paris"}
}
- system messages collected and joined (preserves guardrail prompts)
- assistant messages with tool_calls tool_use content blocks
- tool role messages user role with tool_result content blocks
"""
system_parts: List[str] = []
conversation: List[Dict] = []
for msg in messages:
if isinstance(msg, dict):
role = msg.get("role", "")
content: Any = msg.get("content", "")
else:
role = getattr(msg, "role", "")
content = getattr(msg, "content", "")
if role == "system":
if isinstance(content, str) and content:
system_parts.append(content)
elif isinstance(content, list):
system_parts.append(
"\n".join(
b.get("text", "")
for b in content
if b.get("type") == "text"
)
)
elif role == "assistant":
tool_calls = (
msg.get("tool_calls")
if isinstance(msg, dict)
else getattr(msg, "tool_calls", None)
)
if tool_calls: # type: ignore[truthy-bool]
content_blocks: List[Dict[str, Any]] = []
if content:
content_blocks.append({"type": "text", "text": content})
for tc in tool_calls: # type: ignore[attr-defined]
func = (
tc.get("function", {})
if isinstance(tc, dict)
else getattr(tc, "function", {})
)
tc_id = (
tc.get("id", "")
if isinstance(tc, dict)
else getattr(tc, "id", "")
)
func_name = (
func.get("name", "")
if isinstance(func, dict)
else getattr(func, "name", "")
)
func_args = (
func.get("arguments", "{}")
if isinstance(func, dict)
else getattr(func, "arguments", "{}")
)
try:
input_data = (
json.loads(func_args)
if isinstance(func_args, str)
else func_args
)
except (json.JSONDecodeError, TypeError):
input_data = {}
content_blocks.append(
{
"type": "tool_use",
"id": tc_id,
"name": func_name,
"input": input_data,
}
)
conversation.append(
{"role": "assistant", "content": content_blocks}
)
else:
conversation.append({"role": "assistant", "content": content})
elif role == "tool":
tool_call_id = (
msg.get("tool_call_id", "")
if isinstance(msg, dict)
else getattr(msg, "tool_call_id", "")
)
tool_content = (
content if isinstance(content, str) else json.dumps(content)
)
tool_result_block = {
"type": "tool_result",
"tool_use_id": tool_call_id,
"content": tool_content,
}
if (
conversation
and conversation[-1]["role"] == "user"
and isinstance(conversation[-1]["content"], list)
and conversation[-1]["content"]
and conversation[-1]["content"][0].get("type") == "tool_result"
):
conversation[-1]["content"].append(tool_result_block)
else:
conversation.append(
{"role": "user", "content": [tool_result_block]}
)
else:
conversation.append({"role": role, "content": content})
system: Optional[str] = "\n\n".join(system_parts) if system_parts else None
return system, conversation
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
stream: bool = optional_params.pop("stream", False) or False
extra_body = optional_params.pop("extra_body", {})
if _is_claude_model(model):
return self._transform_request_anthropic(
model, messages, optional_params, stream, extra_body
)
return self._transform_request_openai(
model, messages, optional_params, stream, extra_body
)
def _transform_request_openai(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
stream: bool,
extra_body: dict,
) -> dict:
"""OpenAI format for /chat/completions endpoint."""
max_tokens = optional_params.pop("max_tokens", None)
max_completion_tokens = optional_params.pop("max_completion_tokens", None)
resolved_max = max_completion_tokens or max_tokens
body: dict = {
"model": model.removeprefix("snowflake/"),
"messages": messages,
"stream": stream,
**optional_params,
**extra_body,
}
OpenAI format (returned tool_calls):
ChatCompletionMessageToolCall(
id="tooluse_...",
type="function",
function=Function(name="get_weather", arguments='{"location": "Paris"}')
)
if resolved_max is not None:
body["max_completion_tokens"] = resolved_max
return body
def _transform_tool_choice_to_anthropic(self, tool_choice: Any) -> Dict[str, Any]:
"""
text_content = ""
tool_calls: List[ChatCompletionMessageToolCall] = []
Convert tool_choice from OpenAI format to Anthropic format.
for idx, content_item in enumerate(content_list):
if content_item.get("type") == "text":
text_content += content_item.get("text", "")
OpenAI string values: "auto", "required", "none"
OpenAI dict: {"type": "function", "function": {"name": "..."}}
Anthropic: {"type": "auto"}, {"type": "any"}, {"type": "tool", "name": "..."}
"""
if isinstance(tool_choice, str):
mapping = {
"auto": {"type": "auto"},
"required": {"type": "any"},
"none": {"type": "none"},
}
return mapping.get(tool_choice, {"type": "auto"})
elif isinstance(tool_choice, dict):
if tool_choice.get("type") == "function":
func = tool_choice.get("function", {})
return {"type": "tool", "name": func.get("name", "")}
return tool_choice
return {"type": "auto"}
## TOOL CALLING
elif content_item.get("type") == "tool_use":
tool_use_data = content_item.get("tool_use", {})
tool_call = ChatCompletionMessageToolCall(
id=tool_use_data.get("tool_use_id", ""),
type="function",
function=Function(
name=tool_use_data.get("name", ""),
arguments=json.dumps(tool_use_data.get("input", {})),
),
)
tool_calls.append(tool_call)
def _transform_request_anthropic(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
stream: bool,
extra_body: dict,
) -> dict:
"""Anthropic Messages format for /messages endpoint."""
system, conversation = self._extract_system_and_messages(messages)
return text_content, tool_calls if tool_calls else None
if "tools" in optional_params:
optional_params["tools"] = self._transform_tools_to_anthropic(
optional_params["tools"]
)
if "tool_choice" in optional_params:
optional_params["tool_choice"] = self._transform_tool_choice_to_anthropic(
optional_params["tool_choice"]
)
max_completion_tokens = optional_params.pop("max_completion_tokens", None)
if max_completion_tokens and "max_tokens" not in optional_params:
optional_params["max_tokens"] = max_completion_tokens
model_name = model.removeprefix("snowflake/")
body: Dict[str, Any] = {
"model": model_name,
"messages": conversation,
"stream": stream,
**optional_params,
**extra_body,
}
if system is not None:
body["system"] = system
if "max_tokens" not in body:
body["max_tokens"] = (
4096 # reasonable default; Anthropic API max varies by model
)
return body
def transform_response(
self,
@ -99,6 +388,24 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ModelResponse:
if _is_claude_model(model):
return self._transform_response_anthropic(
model, raw_response, model_response, logging_obj, request_data, messages
)
return self._transform_response_openai(
model, raw_response, model_response, logging_obj, request_data, messages
)
def _transform_response_openai(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: List[AllMessageValues],
) -> ModelResponse:
"""Parse standard OpenAI chat completions response."""
response_json = raw_response.json()
logging_obj.post_call(
@ -108,180 +415,278 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
additional_args={"complete_input_dict": request_data},
)
## RESPONSE TRANSFORMATION
# Snowflake returns content_list (not content) with tool_use objects
# We need to transform this to OpenAI's format with content + tool_calls
if "choices" in response_json and len(response_json["choices"]) > 0:
choice = response_json["choices"][0]
if "message" in choice and "content_list" in choice["message"]:
content_list = choice["message"]["content_list"]
(
text_content,
tool_calls,
) = self._transform_tool_calls_from_snowflake_to_openai(content_list)
# Update the choice message with OpenAI format
choice["message"]["content"] = text_content
if tool_calls:
choice["message"]["tool_calls"] = tool_calls
# Remove Snowflake-specific content_list
del choice["message"]["content_list"]
returned_response = ModelResponse(**response_json)
returned_response.model = "snowflake/" + (returned_response.model or "")
if model is not None:
returned_response._hidden_params["model"] = model
return returned_response
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
If api_base is not provided, use the default DeepSeek /chat/completions endpoint.
"""
api_base = self._get_api_base(api_base, optional_params)
return f"{api_base}/cortex/inference:complete"
def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Transform OpenAI tool format to Snowflake tool format.
Args:
tools: List of tools in OpenAI format
Returns:
List of tools in Snowflake format
OpenAI format:
{
"type": "function",
"function": {
"name": "get_weather",
"description": "...",
"parameters": {...}
}
}
Snowflake format:
{
"tool_spec": {
"type": "generic",
"name": "get_weather",
"description": "...",
"input_schema": {...}
}
}
"""
snowflake_tools: List[Dict[str, Any]] = []
for tool in tools:
if tool.get("type") == "function":
function = tool.get("function", {})
snowflake_tool: Dict[str, Any] = {
"tool_spec": {
"type": "generic",
"name": function.get("name"),
"input_schema": function.get(
"parameters",
{"type": "object", "properties": {}},
),
}
}
# Add description if present
if "description" in function:
snowflake_tool["tool_spec"]["description"] = function["description"]
snowflake_tools.append(snowflake_tool)
return snowflake_tools
def _transform_tool_choice(
self, tool_choice: Union[str, Dict[str, Any]]
) -> Dict[str, Any]:
"""
Transform OpenAI tool_choice format to Snowflake format.
Snowflake requires tool_choice to be an object, not a string.
Ref: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/reference/cortex-inference#post--api-v2-cortex-inference-complete-req-body-schema
Args:
tool_choice: Tool choice in OpenAI format (str or dict)
Returns:
Tool choice in Snowflake format (always an object, never a string)
OpenAI format (string):
"auto", "required", "none"
OpenAI format (dict):
{"type": "function", "function": {"name": "get_weather"}}
Snowflake format:
{"type": "auto"} / {"type": "any"} / {"type": "none"}
{"type": "tool", "name": ["get_weather"]}
Snowflake's API (like Anthropic) requires tool_choice as an object
with a "type" field, not as a bare string.
"""
if isinstance(tool_choice, str):
# Snowflake requires object format, not string.
# Map OpenAI string values to Snowflake object format.
# "required" maps to "any" (Snowflake/Anthropic convention).
_type_map = {
"auto": "auto",
"required": "any",
"none": "none",
}
mapped_type = _type_map.get(tool_choice, tool_choice)
return {"type": mapped_type}
if isinstance(tool_choice, dict):
if tool_choice.get("type") == "function":
function_name = tool_choice.get("function", {}).get("name")
if function_name:
return {
"type": "tool",
"name": [function_name], # Snowflake expects array
}
return tool_choice
def transform_request(
def _transform_response_anthropic(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
stream: bool = optional_params.pop("stream", None) or False
extra_body = optional_params.pop("extra_body", {})
) -> ModelResponse:
"""Parse Anthropic Messages response into OpenAI format."""
response_json = raw_response.json()
## TOOL CALLING
# Transform tools from OpenAI format to Snowflake's tool_spec format
tools = optional_params.pop("tools", None)
if tools:
optional_params["tools"] = self._transform_tools(tools)
logging_obj.post_call(
input=messages,
api_key="",
original_response=response_json,
additional_args={"complete_input_dict": request_data},
)
# Transform tool_choice from OpenAI format to Snowflake's tool name array format
tool_choice = optional_params.pop("tool_choice", None)
if tool_choice:
optional_params["tool_choice"] = self._transform_tool_choice(tool_choice)
text_content = ""
tool_calls = []
return {
"model": model,
"messages": messages,
"stream": stream,
**optional_params,
**extra_body,
for block in response_json.get("content", []):
if block.get("type") == "text":
text_content += block.get("text", "")
elif block.get("type") == "tool_use":
tool_calls.append(
ChatCompletionMessageToolCall(
id=block.get("id", ""),
type="function",
function=Function(
name=block.get("name", ""),
arguments=json.dumps(block.get("input", {})),
),
)
)
_stop_reason_map = {
"end_turn": "stop",
"max_tokens": "length",
"tool_use": "tool_calls",
"stop_sequence": "stop",
}
finish_reason = _stop_reason_map.get(
response_json.get("stop_reason", "end_turn"), "stop"
)
message = Message(content=text_content or None, role="assistant")
if tool_calls:
message.tool_calls = tool_calls
choice = Choices(
finish_reason=finish_reason,
index=0,
message=message,
)
usage_data = response_json.get("usage", {})
usage = Usage(
prompt_tokens=usage_data.get("input_tokens", 0),
completion_tokens=usage_data.get("output_tokens", 0),
total_tokens=usage_data.get("input_tokens", 0)
+ usage_data.get("output_tokens", 0),
)
model_response.choices = [choice]
model_response.usage = usage # type: ignore[attr-defined]
model_response.model = "snowflake/" + response_json.get("model", model)
model_response.id = response_json.get("id", "")
if model is not None:
model_response._hidden_params["model"] = model
return model_response
def get_model_response_iterator(
self,
streaming_response: Any,
sync_stream: bool,
json_mode: Optional[bool] = False,
) -> Any:
return SnowflakeStreamingHandler(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
class SnowflakeStreamingHandler(BaseModelResponseIterator):
"""
Parse streaming events from both Snowflake endpoints.
- /chat/completions: OpenAI SSE format (has "choices" key)
- /messages: Anthropic SSE format (has "type" key like content_block_delta)
"""
def __init__(
self,
streaming_response: Any,
sync_stream: bool,
json_mode: Optional[bool] = False,
):
super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
self._tool_index = 0
self._tool_id = ""
self._tool_name = ""
self._input_tokens = 0
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
if "choices" in chunk:
return self._parse_openai_chunk(chunk)
return self._parse_anthropic_chunk(chunk)
def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk:
choices = chunk.get("choices", [])
if not choices:
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
choice = choices[0]
delta = choice.get("delta", {})
finish_reason = choice.get("finish_reason") or ""
text = delta.get("content") or ""
tool_use = None
tool_calls = delta.get("tool_calls")
if tool_calls:
tc = tool_calls[0]
func = tc.get("function", {})
tool_use = ChatCompletionToolCallChunk(
id=tc.get("id", ""),
type="function",
function={
"name": func.get("name", ""),
"arguments": func.get("arguments", ""),
},
index=tc.get("index", 0),
)
return GenericStreamingChunk(
text=text,
is_finished=finish_reason != "",
finish_reason=finish_reason,
usage=None,
index=choice.get("index", 0),
tool_use=tool_use,
)
def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk:
event_type = chunk.get("type", "")
if event_type == "message_start":
message = chunk.get("message", {})
usage_data = message.get("usage", {})
self._input_tokens = usage_data.get("input_tokens", 0)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
elif event_type == "content_block_delta":
delta = chunk.get("delta", {})
delta_type = delta.get("type", "")
if delta_type == "text_delta":
return GenericStreamingChunk(
text=delta.get("text", ""),
is_finished=False,
finish_reason="",
usage=None,
index=chunk.get("index", 0),
tool_use=None,
)
elif delta_type == "input_json_delta":
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=chunk.get("index", 0),
tool_use=ChatCompletionToolCallChunk(
id=self._tool_id,
type="function",
function={
"name": self._tool_name,
"arguments": delta.get("partial_json", ""),
},
index=self._tool_index,
),
)
elif event_type == "content_block_start":
content_block = chunk.get("content_block", {})
if content_block.get("type") == "tool_use":
self._tool_id = content_block.get("id", "")
self._tool_name = content_block.get("name", "")
self._tool_index = chunk.get("index", 0)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=chunk.get("index", 0),
tool_use=ChatCompletionToolCallChunk(
id=self._tool_id,
type="function",
function={"name": self._tool_name, "arguments": ""},
index=self._tool_index,
),
)
elif event_type == "message_delta":
delta = chunk.get("delta", {})
stop_reason = delta.get("stop_reason", "")
usage_data = chunk.get("usage", {})
_stop_map = {
"end_turn": "stop",
"max_tokens": "length",
"tool_use": "tool_calls",
"stop_sequence": "stop",
}
usage = None
if usage_data or self._input_tokens:
output_t = usage_data.get("output_tokens", 0)
input_t = self._input_tokens or usage_data.get("input_tokens", 0)
usage = ChatCompletionUsageBlock(
prompt_tokens=input_t,
completion_tokens=output_t,
total_tokens=input_t + output_t,
)
return GenericStreamingChunk(
text="",
is_finished=True,
finish_reason=_stop_map.get(stop_reason, "stop"),
usage=usage,
index=0,
tool_use=None,
)
elif event_type == "message_stop":
return GenericStreamingChunk(
text="",
is_finished=True,
finish_reason="stop",
usage=None,
index=0,
tool_use=None,
)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)

View file

@ -0,0 +1,183 @@
"""
Transform request/response for Voyage multimodal embeddings.
Voyage multimodal models use /v1/multimodalembeddings and accept `inputs`
containing content blocks, unlike standard Voyage embeddings which use
/v1/embeddings and a string/list `input` field.
"""
from typing import Any, Dict, List, Optional, Union
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
class VoyageMultimodalEmbeddingError(BaseLLMException):
def __init__(
self,
status_code: int,
message: str,
headers: Union[dict, httpx.Headers] = {},
):
self.status_code = status_code
self.message = message
self.request = httpx.Request(
method="POST", url="https://api.voyageai.com/v1/multimodalembeddings"
)
self.response = httpx.Response(status_code=status_code, request=self.request)
super().__init__(
status_code=status_code,
message=message,
headers=headers,
)
class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig):
"""
Reference: https://docs.voyageai.com/reference/multimodal-embeddings-api
"""
@staticmethod
def is_multimodal_embeddings(model: str) -> bool:
return "multimodal" in model.lower()
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
if api_base:
if not api_base.endswith("/multimodalembeddings"):
api_base = f"{api_base}/multimodalembeddings"
return api_base
return "https://api.voyageai.com/v1/multimodalembeddings"
def get_supported_openai_params(self, model: str) -> list:
return ["dimensions"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
if "dimensions" in non_default_params:
optional_params["output_dimension"] = non_default_params["dimensions"]
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
if api_key is None:
api_key = (
get_secret_str("VOYAGE_API_KEY")
or get_secret_str("VOYAGE_AI_API_KEY")
or get_secret_str("VOYAGE_AI_TOKEN")
)
if not api_key:
raise ValueError(
"Voyage API key is required for multimodal embeddings. "
"Set VOYAGE_API_KEY / VOYAGE_AI_API_KEY / VOYAGE_AI_TOKEN "
"or pass `api_key` explicitly."
)
return {"Authorization": f"Bearer {api_key}"}
def _normalize_content_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
item_type = item.get("type")
if item_type == "image_url":
image_url = item.get("image_url")
if isinstance(image_url, dict):
image_url = image_url.get("url")
if image_url is None:
raise ValueError(
"Voyage multimodal embeddings require a non-empty `image_url`. "
"Got an image content block without a `url`."
)
if isinstance(image_url, str) and image_url.startswith("data:image/"):
_, _, encoded = image_url.partition(",")
return {"type": "image_base64", "image_base64": encoded}
return {"type": "image_url", "image_url": image_url}
return item
def _normalize_input_item(self, item: Any) -> Dict[str, Any]:
if isinstance(item, str):
return {"content": [{"type": "text", "text": item}]}
if isinstance(item, dict) and "content" in item:
content = item.get("content") or []
return {
**item,
"content": [
self._normalize_content_item(content_item)
for content_item in content
],
}
return item
def transform_embedding_request(
self,
model: str,
input: AllEmbeddingInputValues,
optional_params: dict,
headers: dict,
) -> dict:
inputs = input if isinstance(input, list) else [input]
return {
"inputs": [self._normalize_input_item(item) for item in inputs],
"model": model,
**optional_params,
}
def transform_embedding_response(
self,
model: str,
raw_response: httpx.Response,
model_response: EmbeddingResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
request_data: dict = {},
optional_params: dict = {},
litellm_params: dict = {},
) -> EmbeddingResponse:
try:
raw_response_json = raw_response.json()
except Exception:
raise VoyageMultimodalEmbeddingError(
message=raw_response.text, status_code=raw_response.status_code
)
model_response.model = raw_response_json.get("model")
model_response.data = raw_response_json.get("data")
model_response.object = raw_response_json.get("object")
usage_payload = raw_response_json.get("usage", {})
total_tokens = usage_payload.get("total_tokens", 0)
model_response.usage = Usage(
prompt_tokens=total_tokens,
total_tokens=total_tokens,
)
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return VoyageMultimodalEmbeddingError(
message=error_message, status_code=status_code, headers=headers
)

View file

@ -35852,7 +35852,17 @@
"max_input_tokens": 32000,
"max_tokens": 32000,
"mode": "embedding",
"output_cost_per_token": 0.0
"output_cost_per_token": 0.0,
"supports_vision": true
},
"voyage/voyage-multimodal-3.5": {
"input_cost_per_token": 1.2e-07,
"litellm_provider": "voyage",
"max_input_tokens": 32000,
"max_tokens": 32000,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_vision": true
},
"wandb/openai/gpt-oss-120b": {
"max_tokens": 131072,
@ -41753,6 +41763,48 @@
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/google.gemma-4-31b": {
"input_cost_per_token": 1.4e-07,
"output_cost_per_token": 4e-07,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/google.gemma-4-26b-a4b": {
"input_cost_per_token": 1.3e-07,
"output_cost_per_token": 4e-07,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/google.gemma-4-e2b": {
"input_cost_per_token": 4e-08,
"output_cost_per_token": 8e-08,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"volcengine/doubao-seed-2-0-pro-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,

View file

@ -2226,6 +2226,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="max response size in MB, if a response is larger than this size it will be rejected",
)
cancel_on_disconnect: Optional[bool] = Field(
None,
description="cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure",
)
infer_model_from_keys: Optional[bool] = Field(
None,
description="for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY)",

View file

@ -1,6 +1,7 @@
import asyncio
import json
import logging
import math
import time
import traceback
from datetime import datetime
@ -49,6 +50,7 @@ from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.utils import ServerToolUse
# Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format)
@ -556,6 +558,64 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
return False
_CLIENT_DISCONNECT_DETAIL = "Client disconnected the request"
def _log_llm_api_exception(e: Exception) -> None:
if (
getattr(e, "status_code", None) == 499
and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL
):
verbose_proxy_logger.info(
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled"
)
return
verbose_proxy_logger.exception(
f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}"
)
async def _cancel_llm_call_on_client_disconnect(
request: Request,
llm_api_call: "asyncio.Future[Any]",
disconnect_event: asyncio.Event,
) -> None:
try:
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
disconnect_event.set()
llm_api_call.cancel()
return
except Exception as exc:
verbose_proxy_logger.warning(
"cancel_on_disconnect: request.receive() raised %s; "
"upstream LLM call will not be cancelled on disconnect",
exc,
)
async def _await_llm_call_cancelling_on_disconnect(
request: Request,
llm_api_call: "asyncio.Future[Any]",
) -> Any:
disconnect_event = asyncio.Event()
monitor = asyncio.create_task(
_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event)
)
try:
return await llm_api_call
except asyncio.CancelledError:
if disconnect_event.is_set():
raise HTTPException(
status_code=499,
detail=_CLIENT_DISCONNECT_DETAIL,
)
raise
finally:
monitor.cancel()
class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict):
self.data = data
@ -1224,7 +1284,12 @@ class ProxyBaseLLMRequestProcessing:
*tasks
) # run the moderation check in parallel to the actual llm api call
responses = await llm_responses
if general_settings.get("cancel_on_disconnect", False):
responses = await _await_llm_call_cancelling_on_disconnect(
request, llm_responses
)
else:
responses = await llm_responses
response = responses[1]
@ -2067,6 +2132,10 @@ class ProxyBaseLLMRequestProcessing:
e,
)
def _apply_router_cooldown_retry_after(self, headers: dict, e: Exception) -> None:
if isinstance(e, RouterRateLimitError) and e.cooldown_time > 0:
headers["retry-after"] = str(math.ceil(e.cooldown_time))
async def _handle_llm_api_exception(
self,
e: Exception,
@ -2075,9 +2144,7 @@ class ProxyBaseLLMRequestProcessing:
version: Optional[str] = None,
):
"""Raises ProxyException (OpenAI API compatible) if an exception is raised"""
verbose_proxy_logger.exception(
f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}"
)
_log_llm_api_exception(e)
# Allow callbacks to transform the error response
transformed_exception = await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
@ -2148,6 +2215,8 @@ class ProxyBaseLLMRequestProcessing:
except Exception:
pass
self._apply_router_cooldown_retry_after(headers, e)
if isinstance(e, HTTPException):
raw_detail = getattr(e, "detail", str(e))
message, structured_fields = _serialize_http_exception_detail(raw_detail)

View file

@ -24,6 +24,7 @@ from litellm.proxy._types import (
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
SpecialModelNames,
UserAPIKeyAuth,
WebhookEvent,
)
@ -1074,8 +1075,26 @@ async def health_endpoint(
# response but NOT in the background-cache /health response. This is
# surfaced via the "warnings" field below so operators can fix the
# missing model_info.id rather than guess at the discrepancy.
if len(user_api_key_dict.models) > 0:
allowed_models = set(user_api_key_dict.models)
# Keys granted SpecialModelNames.all_proxy_models carry the literal
# "all-proxy-models" entry, which matches no real model_name; treat
# them as unrestricted instead of filtering the list down to nothing.
# Keys granted SpecialModelNames.all_team_models inherit the parent
# team's allowlist (same semantics as get_key_models in
# model_checks.py). Without a team_id the sentinel cannot resolve and
# stays in the list, matching nothing; denied rather than
# unrestricted, mirroring _resolve_key_models_for_auth_check.
accessible_models = list(user_api_key_dict.models)
if (
SpecialModelNames.all_team_models.value in accessible_models
and user_api_key_dict.team_id is not None
):
accessible_models = list(user_api_key_dict.team_models)
restrict_to_allowed_models = (
len(accessible_models) > 0
and SpecialModelNames.all_proxy_models.value not in accessible_models
)
if restrict_to_allowed_models:
allowed_models = set(accessible_models)
_llm_model_list = [
m for m in _llm_model_list if m.get("model_name") in allowed_models
]
@ -1087,7 +1106,7 @@ async def health_endpoint(
# other healthy model would still report healthy_count > 0 and
# the targeted-503 path would never fire.
targeted_ids = _resolve_targeted_model_ids(_llm_model_list, model, model_id)
if len(user_api_key_dict.models) > 0:
if restrict_to_allowed_models:
allowed_model_ids = {
(m.get("model_info") or {}).get("id")
for m in _llm_model_list

View file

@ -397,6 +397,32 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str
)
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
extensions and the Agent SDK run through the same CLI and share that prefix."""
return user_agent.startswith("claude-cli/")
def should_auto_drop_params_for_claude_code(
user_agent: str, data: dict, proxy_config: ProxyConfig
) -> bool:
"""drop_params defaults to on for Claude Code so its Anthropic-specific
params (e.g. thinking) don't fail requests routed to non-Anthropic
providers. An explicit drop_params from the caller or in the operator's
``litellm_settings`` always wins over this default."""
if not is_claude_code_user_agent(user_agent):
return False
if "drop_params" in data:
return False
config = getattr(proxy_config, "config", None)
litellm_settings = (
config.get("litellm_settings") if isinstance(config, dict) else None
)
return not (
isinstance(litellm_settings, dict) and "drop_params" in litellm_settings
)
def safe_add_api_version_from_query_params(data: dict, request: Request):
try:
if hasattr(request, "query_params"):
@ -1742,6 +1768,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
user_agent = request.headers["user-agent"]
data[_metadata_variable_name]["user_agent"] = user_agent
if should_auto_drop_params_for_claude_code(user_agent, data, proxy_config):
data["drop_params"] = True
# Merge caller-supplied tags (x-litellm-tags header, data["tags"] root-level)
# into request metadata for tag-based routing and spend attribution.
tags = LiteLLMProxyRequestSetup.add_request_tag_to_metadata(

View file

@ -4794,12 +4794,27 @@ async def reset_key_spend_fn(
proxy_logging_obj=proxy_logging_obj,
)
try:
from litellm.proxy.proxy_server import _invalidate_spend_counter
# Set Redis spend counter to the new value so get_current_spend()
# returns the correct amount immediately instead of the stale pre-reset value.
# We use reset_to (not 0.0) so partial resets are reflected correctly.
from litellm.proxy.proxy_server import spend_counter_cache
await _invalidate_spend_counter(counter_key=f"spend:key:{hashed_api_key}")
except Exception:
pass
_counter_key = f"spend:key:{hashed_api_key}"
spend_counter_cache.in_memory_cache.set_cache(
key=_counter_key, value=reset_to, ttl=60
)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=_counter_key, value=reset_to, ttl=60
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "
"Budget checks may use stale value until counter expires.",
_counter_key,
redis_err,
)
max_budget = updated_key.max_budget
budget_reset_at = updated_key.budget_reset_at

View file

@ -4850,15 +4850,34 @@ async def team_model_add(
detail={"error": "Only proxy admin or team admin can modify team models"},
)
updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models)
# Update team. `include` mirrors the relations the auth path consumes
# off the cached team object so that `_refresh_cached_team` doesn't
# null them out — see object_permission_utils.validate_key_search_tools_against_team
# and the MCP/agent authz paths, which treat a missing object_permission
# as "no team-level restriction".
# Atomic array append with dedup at the database level so concurrent
# BYOK model creates don't overwrite each other's team.models entries.
# When the team currently has models=[] (unrestricted access), the
# CASE expression inserts the 'all-proxy-models' sentinel first.
models_to_add = list(data.models)
await prisma_client.db.execute_raw(
'UPDATE "LiteLLM_TeamTable" '
"SET models = ("
" SELECT ARRAY(SELECT DISTINCT unnest("
" CASE WHEN cardinality(COALESCE(models, ARRAY[]::text[])) = 0 "
" THEN ARRAY['all-proxy-models']::text[] "
" ELSE models "
" END || $1::text[]"
" ))"
") "
"WHERE team_id = $2",
models_to_add,
data.team_id,
)
# Re-fetch via update (write-routed) instead of find_unique (read-routed)
# to avoid returning stale data from a read replica. The models column
# was already set by execute_raw above; this just retrieves the row from
# the writer and lets Prisma bump updated_at.
# `include` mirrors the relations the auth path consumes off the cached
# team object so that `_refresh_cached_team` doesn't null them out.
updated_team = await TeamRepository(prisma_client).table.update(
where={"team_id": data.team_id},
data={"models": updated_models},
data={"updated_at": datetime.now(timezone.utc)},
include={"object_permission": True}, # type: ignore
)

View file

@ -1942,34 +1942,6 @@ db_writer_client: Optional[AsyncHTTPHandler] = None
### logger ###
async def check_request_disconnection(request: Request, llm_api_call_task):
"""
Asynchronously checks if the request is disconnected at regular intervals.
If the request is disconnected
- cancel the litellm.router task
- raises an HTTPException with status code 499 and detail "Client disconnected the request".
Parameters:
- request: Request: The request object to check for disconnection.
Returns:
- None
"""
# only run this function for 10 mins -> if these don't get cancelled -> we don't want the server to have many while loops
start_time = time.time()
while time.time() - start_time < 600:
await asyncio.sleep(1)
if await request.is_disconnected():
# cancel the LLM API Call task if any passed - this is passed from individual providers
# Example OpenAI, Azure, VertexAI etc
llm_api_call_task.cancel()
raise HTTPException(
status_code=499,
detail="Client disconnected the request",
)
def _resolve_typed_dict_type(typ):
"""Resolve the actual TypedDict class from a potentially wrapped type."""
from typing_extensions import _TypedDictMeta # type: ignore
@ -4920,9 +4892,12 @@ class ProxyConfig:
combined_id_list = []
## BASE CASES ##
# if llm_router is None or db_models is empty, return 0
if llm_router is None or len(db_models) == 0:
if llm_router is None:
return 0
# NOTE: db_models may be legitimately empty when all DB models have been deleted.
# Do NOT short-circuit on len(db_models) == 0 — we must still evict any
# DB-sourced deployments that are no longer in the DB. The caller
# (_update_llm_router) already guards against None (transient fetch failure).
## DB MODELS ##
for m in db_models:
@ -5072,6 +5047,15 @@ class ProxyConfig:
)
try:
# new_models is None when _get_models_from_db failed (transient DB error).
# Skip the update entirely so we don't evict valid deployments.
if new_models is None:
verbose_proxy_logger.warning(
"_update_llm_router: DB model fetch returned None (transient failure). "
"Skipping router update to preserve existing deployments."
)
return
models_list: list = new_models if isinstance(new_models, list) else []
if llm_router is None and master_key is not None:
verbose_proxy_logger.debug(f"len new_models: {len(models_list)}")
@ -5774,18 +5758,25 @@ class ProxyConfig:
# Check if the object type is in the list (supports both str and enum values)
return any(str(obj) == object_type_str for obj in supported_db_objects)
async def _get_models_from_db(self, prisma_client: PrismaClient) -> list:
async def _get_models_from_db(self, prisma_client: PrismaClient) -> Optional[list]:
"""
Fetch all model deployments from the DB.
Returns:
- list: the rows (may be empty if no models exist)
- None: signals a DB fetch *failure* callers must not treat this
as "all models deleted" and must not evict existing router deployments.
"""
try:
new_models = await ModelRepository(prisma_client).table.find_many()
return new_models
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format(
str(e)
)
)
new_models = []
return new_models
return None
async def add_deployment(
self,
@ -14775,6 +14766,7 @@ async def get_config_list(
"always_include_stream_usage": {"type": "Boolean"},
"forward_client_headers_to_llm_api": {"type": "Boolean"},
"mcp_required_fields": {"type": "List"},
"cancel_on_disconnect": {"type": "Boolean"},
}
return_val = []

View file

@ -4112,47 +4112,13 @@ class Router:
```
"""
try:
kwargs["model"] = model
kwargs["input"] = input
kwargs["voice"] = voice
deployment = await self.async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
kwargs["original_function"] = self._aspeech
self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs)
data = deployment["litellm_params"].copy()
data["model"]
for k, v in self.default_litellm_params.items():
if (
k not in kwargs
): # prioritize model-specific params > default router params
kwargs[k] = v
elif k == "metadata":
kwargs[k].update(v)
response = await self.async_function_with_fallbacks(**kwargs)
potential_model_client = self._get_client(
deployment=deployment, kwargs=kwargs, client_type="async"
)
# check if provided keys == client keys #
dynamic_api_key = kwargs.get("api_key", None)
if (
dynamic_api_key is not None
and potential_model_client is not None
and dynamic_api_key != potential_model_client.api_key
):
model_client = None
else:
model_client = potential_model_client
response = await litellm.aspeech(
**{
**data,
"client": model_client,
**kwargs,
}
)
return response
except Exception as e:
asyncio.create_task(
@ -4165,6 +4131,76 @@ class Router:
)
raise e
async def _aspeech(self, model: str, input: str, voice: str, **kwargs):
model_name = model
try:
verbose_router_logger.debug(
f"Inside _aspeech()- model: {model}; kwargs: {kwargs}"
)
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
deployment = await self.async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
data = deployment["litellm_params"].copy()
model_client = self._get_async_openai_model_client(
deployment=deployment,
kwargs=kwargs,
)
self.total_calls[model_name] += 1
response = litellm.aspeech(
**{
**data,
"input": input,
"voice": voice,
"client": model_client,
**kwargs,
}
)
### CONCURRENCY-SAFE RPM CHECKS ###
rpm_semaphore = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(
rpm_semaphore, asyncio.Semaphore
):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info(
f"litellm.aspeech(model={model_name})\033[32m 200 OK\033[0m"
)
return response
except Exception as e:
verbose_router_logger.info(
f"litellm.aspeech(model={model_name})\033[31m Exception {str(e)}\033[0m"
)
if model_name is not None:
self.fail_calls[model_name] += 1
raise e
async def arerank(self, model: str, **kwargs):
try:
kwargs["model"] = model

View file

@ -1,4 +1,5 @@
import os
import time
from datetime import datetime as dt
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Set, Union
@ -201,6 +202,8 @@ class HangingRequestData(BaseModel):
key_alias: Optional[str] = None
team_alias: Optional[str] = None
alerting_metadata: Optional[dict] = None
created_at: float = Field(default_factory=time.time)
alerted: bool = False
class AlertTypeConfig(LiteLLMPydanticObjectBase):

View file

@ -3060,6 +3060,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
wandb_api_key: Optional[str]
weave_project_id: Optional[str]
# Datadog dynamic params
dd_api_key: Optional[str]
dd_site: Optional[str]
dd_agent_host: Optional[str]
dd_agent_port: Optional[str]
# Logging settings
turn_off_message_logging: Optional[bool] # when true will not log messages
litellm_disabled_callbacks: Optional[List[str]]

View file

@ -3583,6 +3583,15 @@ def get_optional_params_embeddings( # noqa: PLR0915
drop_params=drop_params if drop_params is not None else False,
)
)
elif litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model):
optional_params = (
litellm.VoyageMultimodalEmbeddingConfig().map_openai_params(
non_default_params=non_default_params,
optional_params={},
model=model,
drop_params=drop_params if drop_params is not None else False,
)
)
else:
optional_params = litellm.VoyageEmbeddingConfig().map_openai_params(
non_default_params=non_default_params,
@ -8666,6 +8675,11 @@ class ProviderConfigManager:
)
):
return litellm.VoyageContextualEmbeddingConfig()
elif (
litellm.LlmProviders.VOYAGE == provider
and litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model)
):
return litellm.VoyageMultimodalEmbeddingConfig()
elif litellm.LlmProviders.VOYAGE == provider:
return litellm.VoyageEmbeddingConfig()
elif litellm.LlmProviders.TRITON == provider:

View file

@ -39511,6 +39511,178 @@
"litellm_provider": "fireworks_ai",
"mode": "chat"
},
"scaleway/qwen/qwen3.5-397b-a17b": {
"input_cost_per_token": 6e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 256000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 3.6e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_vision": true
},
"scaleway/qwen/qwen3.6-35b-a3b": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 256000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"supports_function_calling": true,
"supports_vision": true,
"supports_reasoning": true
},
"scaleway/qwen/qwen3-235b-a22b-instruct-2507": {
"input_cost_per_token": 7.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 256000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 2.25e-06,
"supports_function_calling": true
},
"scaleway/qwen/qwen3-embedding-8b": {
"input_cost_per_token": 1e-07,
"litellm_provider": "scaleway",
"mode": "embedding",
"output_cost_per_token": 0.0
},
"scaleway/qwen/qwen3-coder-30b-a3b-instruct": {
"input_cost_per_token": 2e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 128000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-07,
"supports_function_calling": true
},
"scaleway/openai/gpt-oss-120b": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 128000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 6e-07,
"supports_function_calling": true
},
"scaleway/openai/whisper-large-v3": {
"input_cost_per_audio_token": 0.0,
"litellm_provider": "scaleway",
"mode": "audio_transcription",
"output_cost_per_token": 0.0
},
"scaleway/google/gemma-4-26b-a4b-it": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 256000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 5e-07,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_vision": true
},
"scaleway/google/gemma-3-27b-it": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 40000,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 5e-07,
"supports_function_calling": true,
"supports_vision": true
},
"scaleway/hcompany/holo2-30b-a3b": {
"input_cost_per_token": 3e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 22000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 7e-07,
"supports_reasoning": true,
"supports_vision": true
},
"scaleway/mistralai/mistral-medium-3.5-128b": {
"input_cost_per_token": 1.5e-06,
"litellm_provider": "scaleway",
"max_input_tokens": 256000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 7.5e-06,
"supports_reasoning": true,
"supports_function_calling": true,
"supports_vision": true,
"supports_tool_choice": true
},
"scaleway/mistralai/devstral-2-123b-instruct-2512": {
"input_cost_per_token": 4e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 200000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 2e-06,
"supports_function_calling": true
},
"scaleway/mistralai/voxtral-small-24b-2507": {
"input_cost_per_audio_token": 1.5e-07,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 32000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 3.5e-07,
"supports_audio_input": true
},
"scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 128000,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3.5e-07,
"supports_function_calling": true,
"supports_vision": true
},
"scaleway/mistralai/pixtral-12b-2409": {
"input_cost_per_token": 2e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 128000,
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 2e-07,
"supports_vision": true,
"supports_function_calling": true
},
"scaleway/BAAI/bge-multilingual-gemma2": {
"input_cost_per_token": 1e-07,
"litellm_provider": "scaleway",
"mode": "embedding",
"output_cost_per_token": 0.0
},
"scaleway/meta/llama-3.3-70b-instruct": {
"input_cost_per_token": 9e-07,
"litellm_provider": "scaleway",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 9e-07,
"supports_function_calling": true
},
"novita/deepseek/deepseek-v3.2": {
"litellm_provider": "novita",
"mode": "chat",
@ -41793,6 +41965,48 @@
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/google.gemma-4-31b": {
"input_cost_per_token": 1.4e-07,
"output_cost_per_token": 4e-07,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/google.gemma-4-26b-a4b": {
"input_cost_per_token": 1.3e-07,
"output_cost_per_token": 4e-07,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/google.gemma-4-e2b": {
"input_cost_per_token": 4e-08,
"output_cost_per_token": 8e-08,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 128000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"volcengine/doubao-seed-2-0-pro-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,

View file

@ -2086,7 +2086,7 @@
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": false,
"embeddings": true,
"image_generations": false,
"audio_transcriptions": true,
"audio_speech": false,
@ -2153,7 +2153,7 @@
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": false,
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
@ -2752,6 +2752,23 @@
"batches": false,
"rerank": false
}
},
"empiriolabs": {
"display_name": "EmpirioLabs (`empiriolabs`)",
"url": "https://docs.litellm.ai/docs/providers/empiriolabs",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": false
}
}
},
"endpoints": {

View file

@ -230,6 +230,7 @@ general_settings:
# background_health_checks: true
# use_shared_health_check: true
# health_check_interval: 30
# cancel_on_disconnect: true # cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot)
# database_url: "postgresql://<user>:<password>@<host>:<port>/<dbname>" # [OPTIONAL] use for token-based auth to proxy
pass_through_endpoints:

View file

@ -0,0 +1,63 @@
"""
Unit tests for the EmpirioLabs OpenAI-like provider.
"""
import os
import sys
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
)
from litellm.llms.openai_like.dynamic_config import create_config_class
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
EMPIRIOLABS_BASE_URL = "https://api.empiriolabs.ai/v1"
def _get_config():
provider = JSONProviderRegistry.get("empiriolabs")
assert provider is not None
config_class = create_config_class(provider)
return config_class()
def test_empiriolabs_provider_registered():
provider = JSONProviderRegistry.get("empiriolabs")
assert provider is not None
assert provider.base_url == EMPIRIOLABS_BASE_URL
assert provider.api_key_env == "EMPIRIOLABS_API_KEY"
assert provider.api_base_env == "EMPIRIOLABS_API_BASE"
def test_empiriolabs_resolves_env_api_key(monkeypatch):
config = _get_config()
monkeypatch.setenv("EMPIRIOLABS_API_KEY", "test-key")
api_base, api_key = config._get_openai_compatible_provider_info(None, None)
assert api_base == EMPIRIOLABS_BASE_URL
assert api_key == "test-key"
def test_empiriolabs_maps_max_completion_tokens():
config = _get_config()
params = config.map_openai_params(
non_default_params={"max_completion_tokens": 256},
optional_params={},
model="empiriolabs/qwen3-7-plus",
drop_params=False,
)
assert params.get("max_tokens") == 256
assert "max_completion_tokens" not in params
def test_empiriolabs_complete_url_appends_endpoint():
config = _get_config()
url = config.get_complete_url(
api_base=EMPIRIOLABS_BASE_URL,
api_key="test-key",
model="empiriolabs/qwen3-7-plus",
optional_params={},
litellm_params={},
stream=False,
)
assert url == f"{EMPIRIOLABS_BASE_URL}/chat/completions"

View file

@ -224,8 +224,22 @@ async def test_db_error_new_model_check():
model_info={"id": deployment.model_info.id},
)
db_models = []
deleted_deployments = await pc._delete_deployment(db_models=db_models)
# Mock get_config to return the two deployments as config-backed models so
# they appear in combined_id_list and are not evicted when db_models is empty
# (simulates the real-world case: DB error returns [], but models live in config).
config_model_list = [
deployment.to_json(exclude_none=True),
deployment_2.to_json(exclude_none=True),
]
from unittest.mock import AsyncMock, patch
with patch.object(
pc,
"get_config",
new=AsyncMock(return_value={"model_list": config_model_list}),
):
db_models = []
deleted_deployments = await pc._delete_deployment(db_models=db_models)
assert deleted_deployments == 0
assert init_len_list == len(llm_router.model_list)

View file

@ -198,6 +198,96 @@ async def test_audio_speech_router(mode):
assert test_logger.standard_logging_object["model_group"] == "tts"
@pytest.mark.asyncio
async def test_aspeech_fallbacks_on_deployment_failure():
router = Router(
model_list=[
{
"model_name": "tts-main",
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
},
{
"model_name": "tts-backup",
"litellm_params": {"model": "openai/tts-1-hd", "api_key": "fake-key"},
},
],
fallbacks=[{"tts-main": ["tts-backup"]}],
num_retries=0,
)
called_models = []
async def mock_aspeech(*args, **kwargs):
called_models.append(kwargs["model"])
if kwargs["model"] == "openai/tts-1":
raise litellm.InternalServerError(
message="deployment down",
llm_provider="openai",
model="tts-1",
)
return MagicMock()
with patch("litellm.aspeech", side_effect=mock_aspeech):
response = await router.aspeech(
model="tts-main",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
)
assert response is not None
assert called_models == ["openai/tts-1", "openai/tts-1-hd"]
@pytest.mark.asyncio
async def test_aspeech_success_returns_response():
router = Router(
model_list=[
{
"model_name": "tts",
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
},
]
)
mock_response = MagicMock()
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
response = await router.aspeech(
model="tts",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
)
assert response is mock_response
mock_aspeech.assert_called_once()
assert mock_aspeech.call_args.kwargs["model"] == "openai/tts-1"
@pytest.mark.asyncio
async def test_aspeech_sets_deployment_metadata():
router = Router(
model_list=[
{
"model_name": "tts",
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
},
]
)
mock_response = MagicMock()
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
response = await router._aspeech(
model="tts",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
)
assert response is mock_response
metadata = mock_aspeech.call_args.kwargs["metadata"]
assert metadata["deployment"] == "openai/tts-1"
assert metadata["deployment_model_name"] == "tts"
assert metadata["model_info"]["id"] is not None
@pytest.mark.asyncio()
async def test_rerank_endpoint(model_list):
from litellm.types.utils import RerankResponse

View file

@ -1,6 +1,7 @@
import json
import os
import sys
import time
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
@ -35,13 +36,13 @@ class TestAlertingHangingRequestCheck:
async def test_init_creates_cache_with_correct_ttl(self, mock_slack_alerting):
"""
Test that initialization creates a hanging request cache with correct TTL.
The TTL should be alerting_threshold + buffer time.
The TTL should be 1.5x alerting_threshold + buffer time, so entries
survive long enough to be checked after crossing the threshold.
"""
checker = AlertingHangingRequestCheck(slack_alerting_object=mock_slack_alerting)
# The cache should be created with TTL = alerting_threshold + buffer time
expected_ttl = (
mock_slack_alerting.alerting_threshold + 60
expected_ttl = int(
mock_slack_alerting.alerting_threshold * 1.5 + 60
) # HANGING_ALERT_BUFFER_TIME_SECONDS
assert checker.hanging_request_cache.default_ttl == expected_ttl
@ -208,13 +209,14 @@ class TestAlertingHangingRequestCheck:
Test send_alerts_for_hanging_requests when request is actually hanging.
Should send alert for requests that haven't completed within threshold.
"""
# Add a hanging request to the cache
# Add a hanging request that is older than the alerting threshold
hanging_data = HangingRequestData(
request_id="hanging_request_999",
model="gpt-4",
api_base="https://api.openai.com/v1",
key_alias="test_key",
team_alias="test_team",
created_at=time.time() - 301,
)
await hanging_request_checker.hanging_request_cache.async_set_cache(
key="hanging_request_999", value=hanging_data, ttl=300
@ -236,6 +238,82 @@ class TestAlertingHangingRequestCheck:
# Verify alert was sent for hanging request
hanging_request_checker.slack_alerting_object.send_alert.assert_called_once()
@pytest.mark.asyncio
async def test_send_alerts_for_hanging_requests_alerts_once_per_hang(
self, hanging_request_checker
):
"""
A single hanging request must alert exactly once even though the
checker tick revisits it on every run within the cache TTL.
"""
hanging_data = HangingRequestData(
request_id="hanging_once_555",
model="gpt-4",
api_base="https://api.openai.com/v1",
created_at=time.time() - 301,
)
await hanging_request_checker.hanging_request_cache.async_set_cache(
key="hanging_once_555", value=hanging_data, ttl=300
)
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy:
mock_internal_cache = AsyncMock()
mock_internal_cache.async_get_cache.return_value = None
mock_proxy.internal_usage_cache = mock_internal_cache
hanging_request_checker.hanging_request_cache.async_get_oldest_n_keys = (
AsyncMock(return_value=["hanging_once_555"])
)
for _ in range(3):
await hanging_request_checker.send_alerts_for_hanging_requests()
assert hanging_request_checker.slack_alerting_object.send_alert.call_count == 1
cached = await hanging_request_checker.hanging_request_cache.async_get_cache(
key="hanging_once_555"
)
assert cached is not None
assert cached.alerted is True
@pytest.mark.asyncio
async def test_send_alerts_for_hanging_requests_skips_request_younger_than_threshold(
self, hanging_request_checker
):
"""
Test that an in-flight request younger than the alerting threshold
does not trigger an alert and stays in the cache for later checks.
"""
hanging_data = HangingRequestData(
request_id="young_request_123",
model="gpt-4",
api_base="https://api.openai.com/v1",
)
await hanging_request_checker.hanging_request_cache.async_set_cache(
key="young_request_123", value=hanging_data, ttl=300
)
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy:
# Mock internal usage cache to return None (request still in flight)
mock_internal_cache = AsyncMock()
mock_internal_cache.async_get_cache.return_value = None
mock_proxy.internal_usage_cache = mock_internal_cache
hanging_request_checker.hanging_request_cache.async_get_oldest_n_keys = (
AsyncMock(return_value=["young_request_123"])
)
await hanging_request_checker.send_alerts_for_hanging_requests()
# No alert for a request below the threshold, and it must remain
# cached so a later check can alert if it never completes
hanging_request_checker.slack_alerting_object.send_alert.assert_not_called()
assert (
await hanging_request_checker.hanging_request_cache.async_get_cache(
key="young_request_123"
)
is not None
)
@pytest.mark.asyncio
async def test_send_alerts_for_hanging_requests_with_missing_hanging_data(
self, hanging_request_checker

View file

@ -0,0 +1,263 @@
"""
Tests for team-scoped Datadog callback support.
Verifies that DataDogLogger can be instantiated with per-team credentials
(dd_api_key, dd_site) instead of relying solely on environment variables,
and that the DataDogHandler correctly resolves and caches per-team loggers.
"""
from unittest.mock import patch
import pytest
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.integrations.datadog.datadog_team_handler import (
DataDogHandler,
DatadogLoggingConfig,
)
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
DynamicLoggingCache,
)
from litellm.types.utils import StandardCallbackDynamicParams
@pytest.fixture
def datadog_env(monkeypatch):
"""Set global DD env vars for the default/global logger."""
monkeypatch.setenv("DD_API_KEY", "global_api_key")
monkeypatch.setenv("DD_SITE", "us1.datadoghq.com")
class TestDataDogLoggerCredentialKwargs:
"""Test that DataDogLogger accepts credentials as kwargs."""
def test_init_with_explicit_credentials(self):
"""Logger should use explicit kwargs instead of env vars."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_api_key="team_api_key",
dd_site="eu1.datadoghq.com",
)
assert logger.DD_API_KEY == "team_api_key"
assert "eu1.datadoghq.com" in logger.intake_url
def test_init_falls_back_to_env_vars(self, datadog_env):
"""Logger should fall back to env vars when no kwargs provided."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
assert logger.DD_API_KEY == "global_api_key"
assert "us1.datadoghq.com" in logger.intake_url
def test_init_kwargs_override_env_vars(self, datadog_env):
"""Explicit kwargs should take precedence over env vars."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_api_key="override_key",
dd_site="ap1.datadoghq.com",
)
assert logger.DD_API_KEY == "override_key"
assert "ap1.datadoghq.com" in logger.intake_url
def test_init_with_agent_credentials(self):
"""Logger should use agent mode when dd_agent_host is provided."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_agent_host="dd-agent.local",
dd_agent_port="8125",
dd_api_key="agent_api_key",
)
assert "dd-agent.local:8125" in logger.intake_url
assert logger.DD_API_KEY == "agent_api_key"
def test_init_raises_without_credentials(self, monkeypatch):
"""Logger should raise if no credentials are available."""
monkeypatch.delenv("DD_API_KEY", raising=False)
monkeypatch.delenv("DD_SITE", raising=False)
monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False)
with pytest.raises(Exception, match="DD_API_KEY"):
with patch("asyncio.create_task"):
DataDogLogger()
def test_agent_mode_does_not_leak_env_api_key_when_disallowed(self, datadog_env):
"""With allow_env_credentials=False, the agent logger must not pick up DD_API_KEY env var."""
with patch("asyncio.create_task"):
logger = DataDogLogger(
dd_agent_host="attacker.example.com",
allow_env_credentials=False,
)
assert logger.DD_API_KEY is None
assert "attacker.example.com" in logger.intake_url
def test_direct_api_mode_does_not_leak_env_api_key_when_disallowed(
self, datadog_env
):
"""With allow_env_credentials=False and no explicit key, init must fail rather than reuse env key."""
with pytest.raises(Exception, match="DD_API_KEY"):
with patch("asyncio.create_task"):
DataDogLogger(
dd_site="attacker.example.com",
allow_env_credentials=False,
)
class TestDataDogHandler:
"""Test that DataDogHandler resolves the correct logger per team."""
def test_creates_team_logger_with_dynamic_credentials(self, datadog_env):
"""Should create a new logger when team credentials are provided."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_api_key="team_a_key",
dd_site="eu1.datadoghq.com",
)
with patch("asyncio.create_task"):
result = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result.DD_API_KEY == "team_a_key"
assert "eu1.datadoghq.com" in result.intake_url
def test_caches_team_logger(self, datadog_env):
"""Same team credentials should return the same cached logger instance."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_api_key="team_b_key",
dd_site="us5.datadoghq.com",
)
with patch("asyncio.create_task"):
result1 = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
result2 = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result1 is result2
def test_different_teams_get_different_loggers(self, datadog_env):
"""Different team credentials should create separate logger instances."""
cache = DynamicLoggingCache()
params_a = StandardCallbackDynamicParams(
dd_api_key="team_a_key",
dd_site="us1.datadoghq.com",
)
params_b = StandardCallbackDynamicParams(
dd_api_key="team_b_key",
dd_site="eu1.datadoghq.com",
)
with patch("asyncio.create_task"):
result_a = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params_a,
in_memory_dynamic_logger_cache=cache,
)
result_b = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params_b,
in_memory_dynamic_logger_cache=cache,
)
assert result_a is not result_b
assert result_a.DD_API_KEY == "team_a_key"
assert result_b.DD_API_KEY == "team_b_key"
def test_partial_agent_config_does_not_leak_env_api_key(self, datadog_env):
"""A team-supplied dd_agent_host without dd_api_key must not exfiltrate the proxy DD_API_KEY."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_agent_host="attacker.example.com",
)
with patch("asyncio.create_task"):
result = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result.DD_API_KEY is None
assert "attacker.example.com" in result.intake_url
def test_partial_site_config_does_not_leak_env_api_key(self, datadog_env):
"""A team-supplied dd_site without dd_api_key must not exfiltrate the proxy DD_API_KEY."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_site="attacker.example.com",
)
with pytest.raises(Exception, match="DD_API_KEY"):
with patch("asyncio.create_task"):
DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
def test_full_team_config_still_uses_supplied_key(self, datadog_env):
"""When a team supplies its own key alongside a custom site, that key (not the env key) is used."""
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(
dd_api_key="team_key",
dd_site="eu1.datadoghq.com",
)
with patch("asyncio.create_task"):
result = DataDogHandler.get_datadog_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result.DD_API_KEY == "team_key"
assert "eu1.datadoghq.com" in result.intake_url
def test_request_blocked_callback_params_includes_dd(self):
"""DD params should be blocked from request-level metadata (security)."""
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
_request_blocked_callback_params,
)
assert "dd_api_key" in _request_blocked_callback_params
assert "dd_site" in _request_blocked_callback_params
assert "dd_agent_host" in _request_blocked_callback_params
assert "dd_agent_port" in _request_blocked_callback_params
class TestDynamicCredentialDetection:
"""Test that _dynamic_datadog_credentials_are_passed works correctly."""
def test_no_credentials(self):
params = StandardCallbackDynamicParams()
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False
def test_dd_api_key_only(self):
params = StandardCallbackDynamicParams(dd_api_key="key")
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
def test_dd_site_only(self):
params = StandardCallbackDynamicParams(dd_site="site")
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
def test_dd_agent_host_only(self):
params = StandardCallbackDynamicParams(dd_agent_host="host")
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
class TestStandardCallbackDynamicParamsIncludesDatadog:
"""Verify that Datadog params are in the allow-list."""
def test_dd_params_in_annotations(self):
annotations = StandardCallbackDynamicParams.__annotations__
assert "dd_api_key" in annotations
assert "dd_site" in annotations
assert "dd_agent_host" in annotations
assert "dd_agent_port" in annotations

View file

@ -29,6 +29,7 @@ from litellm.integrations.otel.plumbing.metrics import (
from litellm.integrations.otel.model.payloads import ( # noqa: E402
GuardrailSpanData,
LLMCallSpanData,
LLMCost,
LLMRequestParams,
LLMUsage,
ProxyRequestSpanData,
@ -224,6 +225,97 @@ def test_genai_mapper_all_request_params():
assert attrs["server.port"] == 443
def test_genai_mapper_cost_breakdown():
from litellm.integrations.otel.model.semconv import LiteLLM
data = LLMCallSpanData(
operation=GenAIOperation.CHAT,
provider="anthropic",
request_model="claude-sonnet-4-6",
response_model=None,
response_id=None,
request_params=LLMRequestParams(),
usage=LLMUsage(),
finish_reasons=(),
error=None,
response_cost=0.012,
server=None,
identity=RequestIdentity(call_id=None),
cost=LLMCost(
input=0.004,
output=0.006,
cache_read=0.001,
cache_creation=0.0,
tool_usage=0.0005,
original=0.013,
discount_amount=0.001,
discount_percent=0.077,
margin_total_amount=0.0,
# margin_fixed_amount / margin_percent left unset on purpose
),
)
attrs = GenAIMapper().map(data)
assert attrs[f"{LiteLLM.COST_PREFIX}total"] == 0.012
assert attrs[f"{LiteLLM.COST_PREFIX}input"] == 0.004
assert attrs[f"{LiteLLM.COST_PREFIX}output"] == 0.006
assert attrs[f"{LiteLLM.COST_PREFIX}cache_read"] == 0.001
assert attrs[f"{LiteLLM.COST_PREFIX}cache_creation"] == 0.0
assert attrs[f"{LiteLLM.COST_PREFIX}tool_usage"] == 0.0005
assert attrs[f"{LiteLLM.COST_PREFIX}original"] == 0.013
assert attrs[f"{LiteLLM.COST_PREFIX}discount_amount"] == 0.001
assert attrs[f"{LiteLLM.COST_PREFIX}discount_percent"] == 0.077
assert attrs[f"{LiteLLM.COST_PREFIX}margin_total_amount"] == 0.0
# Components the source did not report are omitted, not zero-filled.
assert f"{LiteLLM.COST_PREFIX}margin_fixed_amount" not in attrs
assert f"{LiteLLM.COST_PREFIX}margin_percent" not in attrs
def test_genai_mapper_cost_breakdown_absent():
# No cost_breakdown → only the rolled-up total (from response_cost) emits.
from litellm.integrations.otel.model.semconv import LiteLLM
attrs = GenAIMapper().map(_full_llm_call())
assert attrs[f"{LiteLLM.COST_PREFIX}total"] == 0.002
assert not any(
k.startswith(LiteLLM.COST_PREFIX) and k != f"{LiteLLM.COST_PREFIX}total"
for k in attrs
)
def test_llm_cost_from_breakdown_maps_costbreakdown_keys():
cost = LLMCost.from_breakdown(
{
"input_cost": 0.004,
"output_cost": 0.006,
"cache_read_cost": 0.001,
"cache_creation_cost": 0.002,
"tool_usage_cost": 0.0005,
"original_cost": 0.013,
"discount_amount": 0.001,
"discount_percent": 0.077,
"margin_fixed_amount": 0.0,
"margin_percent": 0.1,
"margin_total_amount": 0.0011,
"total_cost": 0.012, # carried on response_cost, not LLMCost
}
)
assert cost.input == 0.004
assert cost.output == 0.006
assert cost.cache_read == 0.001
assert cost.cache_creation == 0.002
assert cost.tool_usage == 0.0005
assert cost.original == 0.013
assert cost.discount_amount == 0.001
assert cost.discount_percent == 0.077
assert cost.margin_fixed_amount == 0.0
assert cost.margin_percent == 0.1
assert cost.margin_total_amount == 0.0011
def test_llm_cost_from_breakdown_none_is_empty():
assert LLMCost.from_breakdown(None) == LLMCost()
def test_genai_mapper_guardrail_and_service():
from litellm.integrations.otel.model.semconv import LiteLLM

View file

@ -57,6 +57,42 @@ def _engine(legacy_compat=True):
return SpanEmitter(tracer, cfg), exporter
def test_llm_call_span_cost_breakdown():
engine, exporter = _engine()
data = LLMCallSpanData.from_standard_logging_payload(
_payload(
cost_breakdown={
"input_cost": 0.004,
"output_cost": 0.006,
"cache_read_cost": 0.001,
"total_cost": 0.011,
}
)
)
engine.emit(SpanRole.LLM_CALL, data)
(span,) = exporter.get_finished_spans()
a = span.attributes
# The rolled-up total stays sourced from response_cost.
assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002
# Per-component breakdown now rides the span.
assert a[f"{LiteLLM.COST_PREFIX}input"] == 0.004
assert a[f"{LiteLLM.COST_PREFIX}output"] == 0.006
assert a[f"{LiteLLM.COST_PREFIX}cache_read"] == 0.001
# Unreported components are omitted, not zero-filled.
assert f"{LiteLLM.COST_PREFIX}margin_total_amount" not in a
def test_tracer_scope_carries_litellm_version():
from litellm._version import version as litellm_version
cfg = OpenTelemetryV2Config(exporter="in_memory")
provider, exporter = providers.in_memory_provider(cfg)
tracer = providers.get_tracer(provider, "litellm-test")
tracer.start_span("probe").end()
(span,) = exporter.get_finished_spans()
assert span.instrumentation_scope.version == litellm_version
def test_llm_call_span_golden():
engine, exporter = _engine()
data = LLMCallSpanData.from_standard_logging_payload(_payload())

View file

@ -114,6 +114,9 @@ class TestLangfuseOtelIntegration:
mock_set_attributes.assert_called_once_with(
mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes
)
mock_span.set_attribute.assert_any_call(
"langfuse.observation.type", "generation"
)
def test_set_langfuse_environment_attribute(self):
"""Test that Langfuse environment is set correctly when environment variable is present."""

View file

@ -0,0 +1,41 @@
import json
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
@pytest.mark.parametrize(
"config,model",
[
(AmazonInvokeConfig, "anthropic.claude-3-sonnet-20240229-v1:0"),
(AmazonInvokeConfig, "amazon.titan-text-express-v1"),
(AmazonInvokeConfig, "mistral.mistral-7b-instruct-v0:2"),
(AmazonAnthropicClaudeConfig, "anthropic.claude-sonnet-4-6"),
],
)
def test_transform_request_drops_stream_chunk_size(config, model):
"""stream_chunk_size is a LiteLLM-internal knob for re-chunking the HTTP
response stream. Leaking it into the provider request body makes Bedrock
reject the whole request: ValidationException 'stream_chunk_size: Extra
inputs are not permitted'."""
request_body = config().transform_request(
model=model,
messages=[{"role": "user", "content": "hi"}],
optional_params={"stream": True, "stream_chunk_size": 2048, "max_tokens": 10},
litellm_params={},
headers={},
)
assert "stream_chunk_size" not in json.dumps(request_body)

View file

@ -1,12 +1,21 @@
import os
import sys
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
import litellm
from litellm.llms.bedrock.chat.invoke_handler import (
AWSEventStreamDecoder,
BedrockLLM,
make_call,
make_sync_call,
)
from litellm.llms.custom_httpx.http_handler import HTTPHandler
def test_transform_thinking_blocks_with_redacted_content():
@ -200,3 +209,120 @@ def test_bedrock_converse_streaming_consistent_id():
assert (
response.id == expected_id
), "All chunk IDs must match the one captured from the messageStart event"
@pytest.mark.asyncio
async def test_make_call_does_not_rechunk_stream_by_default():
"""Re-chunking the event stream into fixed 1024-byte blocks holds small
early events (messageStart, contentBlockStart) in httpx's ByteChunker until
1024 bytes accumulate, delaying time-to-first-chunk by the whole generation
when Bedrock trickles bytes (e.g. buffered tool-use streams)."""
response = MagicMock()
response.status_code = 200
client = MagicMock()
client.post = AsyncMock(return_value=response)
await make_call(
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
headers={},
data="{}",
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
)
response.aiter_bytes.assert_called_once_with(chunk_size=None)
@pytest.mark.asyncio
async def test_make_call_honors_explicit_stream_chunk_size():
response = MagicMock()
response.status_code = 200
client = MagicMock()
client.post = AsyncMock(return_value=response)
await make_call(
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
headers={},
data="{}",
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
stream_chunk_size=2048,
)
response.aiter_bytes.assert_called_once_with(chunk_size=2048)
def test_make_sync_call_does_not_rechunk_stream_by_default():
response = MagicMock()
response.status_code = 200
client = MagicMock()
client.post = MagicMock(return_value=response)
make_sync_call(
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
headers={},
data="{}",
signed_json_body=None,
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
)
response.iter_bytes.assert_called_once_with(chunk_size=None)
def test_make_sync_call_honors_explicit_stream_chunk_size():
response = MagicMock()
response.status_code = 200
client = MagicMock()
client.post = MagicMock(return_value=response)
make_sync_call(
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
headers={},
data="{}",
signed_json_body=None,
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
stream_chunk_size=2048,
)
response.iter_bytes.assert_called_once_with(chunk_size=2048)
def test_legacy_bedrock_llm_streaming_does_not_rechunk_by_default():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.iter_bytes = MagicMock(return_value=iter([]))
client = HTTPHandler()
client.post = MagicMock(return_value=mock_response)
BedrockLLM().completion(
model="cohere.command-text-v14",
messages=[{"role": "user", "content": "hi"}],
api_base=None,
custom_prompt_dict={},
model_response=litellm.ModelResponse(),
print_verbose=lambda *args, **kwargs: None,
encoding=litellm.encoding,
logging_obj=MagicMock(),
optional_params={
"stream": True,
"aws_access_key_id": "fake",
"aws_secret_access_key": "fake",
"aws_region_name": "us-east-1",
},
acompletion=False,
timeout=None,
litellm_params={},
client=client,
)
mock_response.iter_bytes.assert_called_once_with(chunk_size=None)

View file

@ -0,0 +1,176 @@
"""
Regression for #30200.
``_auth_with_web_identity_token`` passes an inline ``Policy`` to
``sts.assume_role_with_web_identity``. In AWS IAM an STS session policy
acts as a PERMISSION CEILING effective permissions are the
intersection of the role's identity policies and this policy, so any
action not listed here 403s on OIDC-auth requests only (static creds
and IRSA flow through different paths).
The original policy only granted ``bedrock:*`` actions. When
``#27678`` added the ``bedrock/claude_platform/<model>`` route, the
service-side action namespace was ``aws-external-anthropic:*``, not
``bedrock:*``, so every claude_platform call via OIDC silently denied
with::
User: arn:aws:sts::ACCOUNT:assumed-role/...
is not authorized to perform: aws-external-anthropic:CreateInference
on resource: arn:aws:aws-external-anthropic:...
because no session policy allows the
aws-external-anthropic:CreateInference action
even with a fully permissive identity policy.
Tests below intercept the kwargs handed to
``assume_role_with_web_identity``, parse the embedded ``Policy`` JSON,
and assert that both the original bedrock statement and the new
claude_platform statement are present and cover every documented
action.
"""
import json
from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, patch
import pytest
# Actions the Claude Platform on AWS service is documented to call.
# Source: AWS IAM action reference + the #27678 surface area.
_CLAUDE_PLATFORM_ACTIONS = {
"aws-external-anthropic:CreateInference",
"aws-external-anthropic:CreateBatchInference",
"aws-external-anthropic:CancelBatchInference",
"aws-external-anthropic:DeleteBatchInference",
"aws-external-anthropic:CountTokens",
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
}
def _captured_policy() -> dict:
"""Run _auth_with_web_identity_token under mocks + return the parsed
Policy dict that was actually sent to STS."""
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
base = BaseAWSLLM()
mock_sts = MagicMock()
mock_sts.assume_role_with_web_identity.return_value = {
"Credentials": {
"AccessKeyId": "k",
"SecretAccessKey": "s",
"SessionToken": "t",
"Expiration": datetime.now(timezone.utc) + timedelta(hours=1),
},
"PackedPolicySize": 0,
}
with (
patch("boto3.client", return_value=mock_sts),
patch(
"litellm.llms.bedrock.base_aws_llm.get_secret",
return_value="oidc-jwt-token",
),
):
base._auth_with_web_identity_token(
aws_web_identity_token="/path/to/token",
aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role",
aws_session_name="test-session",
aws_region_name="us-east-1",
aws_sts_endpoint=None,
)
mock_sts.assume_role_with_web_identity.assert_called_once()
kwargs = mock_sts.assume_role_with_web_identity.call_args.kwargs
policy_str = kwargs["Policy"]
return json.loads(policy_str)
def _statement_by_sid(policy: dict, sid: str) -> dict:
for stmt in policy["Statement"]:
if stmt.get("Sid") == sid:
return stmt
raise AssertionError(
f"Sid={sid!r} not found in session policy; "
f"saw {[s.get('Sid') for s in policy['Statement']]}"
)
class TestWebIdentitySessionPolicyShape:
def test_policy_parses_as_valid_iam_document(self):
policy = _captured_policy()
assert policy["Version"] == "2012-10-17"
assert isinstance(policy["Statement"], list)
assert len(policy["Statement"]) >= 2
def test_bedrock_statement_actions_preserved(self):
"""The original bedrock action set must still be granted —
regression for the pre-existing bedrock/* routes."""
policy = _captured_policy()
bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM")
actions = set(bedrock_stmt["Action"])
for required in (
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
):
assert required in actions, f"{required} missing from BedrockLiteLLM"
class TestClaudePlatformActionsCovered:
"""The #30200 bug: every action in the claude_platform service
namespace must appear in the session policy or OIDC requests 403."""
@pytest.mark.parametrize("action", sorted(_CLAUDE_PLATFORM_ACTIONS))
def test_claude_platform_action_present(self, action: str):
policy = _captured_policy()
# Action may live in any Statement — search across all.
all_actions: set = set()
for stmt in policy["Statement"]:
stmt_actions = stmt.get("Action")
if isinstance(stmt_actions, str):
all_actions.add(stmt_actions)
elif isinstance(stmt_actions, list):
all_actions.update(stmt_actions)
assert action in all_actions, (
f"{action} missing from session policy — "
f"bedrock/claude_platform/* requests will 403 on OIDC auth"
)
def test_claude_platform_statement_allows(self):
policy = _captured_policy()
stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM")
assert stmt["Effect"] == "Allow"
assert stmt["Resource"] == "*"
def test_no_aws_external_anthropic_statement_collision(self):
"""Don't accidentally grant a `*` action that would broaden the
ceiling beyond what the documented actions require."""
policy = _captured_policy()
stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM")
actions = stmt["Action"]
if isinstance(actions, str):
actions = [actions]
assert "aws-external-anthropic:*" not in actions, (
"session policy must not grant aws-external-anthropic:* — "
"the ceiling should match the documented action set"
)
class TestPolicyTransportConditions:
def test_bedrock_statement_keeps_secure_transport_condition(self):
policy = _captured_policy()
bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM")
cond = bedrock_stmt.get("Condition") or {}
assert cond.get("Bool", {}).get("aws:SecureTransport") == "true"
def test_claude_platform_statement_carries_secure_transport_condition(self):
"""The new statement should match the existing one's hardening
posture TLS-only, same as bedrock."""
policy = _captured_policy()
stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM")
cond = stmt.get("Condition") or {}
assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", (
"ClaudePlatformLiteLLM must require aws:SecureTransport=true "
"to keep parity with the bedrock statement"
)

View file

@ -20,6 +20,23 @@ from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatCon
from litellm.types.utils import LlmProviders
@pytest.fixture
def local_cost_map(monkeypatch):
original_model_cost = litellm.model_cost
original_bedrock_mantle_models = set(litellm.bedrock_mantle_models)
try:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm.get_model_info.cache_clear()
litellm.add_known_models()
yield
finally:
litellm.model_cost = original_model_cost
litellm.bedrock_mantle_models.clear()
litellm.bedrock_mantle_models.update(original_bedrock_mantle_models)
litellm.get_model_info.cache_clear()
class TestBedrockMantleProviderRegistration:
def test_provider_enum_exists(self):
assert LlmProviders.BEDROCK_MANTLE == "bedrock_mantle"
@ -310,3 +327,52 @@ class TestBedrockMantlePricing:
litellm.add_known_models()
info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
assert info["max_input_tokens"] == 131072
@pytest.mark.parametrize(
"model_id,input_cost,output_cost,max_tokens",
[
("google.gemma-4-31b", 1.4e-07, 4e-07, 256000),
("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000),
("google.gemma-4-e2b", 4e-08, 8e-08, 128000),
],
)
def test_gemma_4_bedrock_mantle_model_metadata(
local_cost_map, model_id, input_cost, output_cost, max_tokens
):
full_model_name = f"bedrock_mantle/{model_id}"
info = litellm.get_model_info(full_model_name)
assert info["mode"] == "chat"
assert info["input_cost_per_token"] == pytest.approx(input_cost)
assert info["output_cost_per_token"] == pytest.approx(output_cost)
assert info["max_input_tokens"] == max_tokens
assert info["max_output_tokens"] == max_tokens
assert info["supports_function_calling"] is True
assert info["supports_reasoning"] is True
assert info["supports_tool_choice"] is True
assert info["supports_vision"] is True
assert (
litellm.supports_parallel_function_calling(
model=full_model_name, custom_llm_provider="bedrock_mantle"
)
is False
)
@pytest.mark.parametrize(
"model_id",
[
"google.gemma-4-31b",
"google.gemma-4-26b-a4b",
"google.gemma-4-e2b",
],
)
def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id):
full_model_name = f"bedrock_mantle/{model_id}"
assert full_model_name in litellm.bedrock_mantle_models
resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name)
assert provider == "bedrock_mantle"
assert resolved_model == model_id

View file

@ -1,10 +1,14 @@
import os
import sys
from unittest.mock import MagicMock
import pytest
import litellm
from litellm.llms.bedrock.chat import BedrockConverseLLM
from litellm.llms.bedrock.chat.converse_handler import make_sync_call
from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions
from litellm.llms.custom_httpx.http_handler import HTTPHandler
sys.path.insert(
0, os.path.abspath("../../../../..")
@ -133,3 +137,79 @@ class TestBedrockRegionInModelPath:
assert model_id == "moonshotai.kimi-k2.5"
# explicitly set region is preserved
assert optional_params["aws_region_name"] == "eu-west-1"
def _stream_completion_with_spied_iter_bytes(model: str, **kwargs) -> MagicMock:
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.iter_bytes = MagicMock(return_value=iter([]))
client = HTTPHandler()
client.post = MagicMock(return_value=mock_response)
litellm.completion(
model=model,
messages=[{"role": "user", "content": "hi"}],
stream=True,
client=client,
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-east-1",
**kwargs,
)
return mock_response.iter_bytes
def test_make_sync_call_does_not_rechunk_stream_by_default():
"""Re-chunking the event stream into fixed 1024-byte blocks holds small
early events in httpx's ByteChunker until 1024 bytes accumulate, delaying
time-to-first-chunk by the whole generation when Bedrock trickles bytes
(e.g. buffered tool-use streams)."""
response = MagicMock()
response.status_code = 200
client = MagicMock()
client.post = MagicMock(return_value=response)
make_sync_call(
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
headers={},
data="{}",
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
)
response.iter_bytes.assert_called_once_with(chunk_size=None)
def test_make_sync_call_honors_explicit_stream_chunk_size():
response = MagicMock()
response.status_code = 200
client = MagicMock()
client.post = MagicMock(return_value=response)
make_sync_call(
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
headers={},
data="{}",
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
stream_chunk_size=2048,
)
response.iter_bytes.assert_called_once_with(chunk_size=2048)
def test_completion_plumbs_stream_chunk_size_through_converse():
iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
)
iter_bytes_spy.assert_called_once_with(chunk_size=None)
iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
stream_chunk_size=2048,
)
iter_bytes_spy.assert_called_once_with(chunk_size=2048)

View file

@ -79,6 +79,20 @@ class TestTensormeshProviderConfig:
matching the text_completion flag in provider_endpoints_support.json."""
assert "tensormesh" in litellm.openai_text_completion_compatible_providers
def test_tensormesh_responses_api_enabled(self):
"""Tensormesh declares /v1/responses in supported_endpoints, so litellm
resolves a responses config for it."""
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.utils import ProviderConfigManager
assert JSONProviderRegistry.supports_responses_api("tensormesh") is True
config = ProviderConfigManager.get_provider_responses_api_config(
provider="tensormesh",
model="tensormesh/openai/gpt-oss-120b",
)
assert config is not None
assert config.custom_llm_provider == "tensormesh"
def test_tensormesh_router_config(self):
"""Test that tensormesh can be used in Router configuration"""
from litellm import Router

View file

@ -26,11 +26,13 @@ class TestSnowflakeToolTransformation:
def test_transform_request_with_tools(self):
"""
Test that OpenAI tool format is correctly transformed to Snowflake's tool_spec format.
Test that OpenAI tool format is passed through as-is to the native endpoint.
The native /chat/completions endpoint accepts standard OpenAI tool format
directly no Snowflake-specific tool_spec transformation needed.
"""
config = SnowflakeConfig()
# OpenAI format tools
tools = [
{
"type": "function",
@ -58,113 +60,94 @@ class TestSnowflakeToolTransformation:
optional_params = {"tools": tools}
transformed_request = config.transform_request(
model="claude-3-5-sonnet",
model="llama3.1-70b",
messages=[{"role": "user", "content": "What's the weather?"}],
optional_params=optional_params,
litellm_params={},
headers={},
)
# Verify tools were transformed to Snowflake format
assert "tools" in transformed_request
assert len(transformed_request["tools"]) == 1
snowflake_tool = transformed_request["tools"][0]
assert "tool_spec" in snowflake_tool
assert snowflake_tool["tool_spec"]["type"] == "generic"
assert snowflake_tool["tool_spec"]["name"] == "get_weather"
assert (
snowflake_tool["tool_spec"]["description"]
== "Get the current weather in a given location"
)
assert "input_schema" in snowflake_tool["tool_spec"]
assert snowflake_tool["tool_spec"]["input_schema"]["type"] == "object"
assert "location" in snowflake_tool["tool_spec"]["input_schema"]["properties"]
assert transformed_request["tools"] == tools
assert "tool_spec" not in json.dumps(transformed_request)
def test_transform_request_with_tool_choice(self):
"""
Test that OpenAI tool_choice format is correctly transformed to Snowflake format.
Test that OpenAI tool_choice format is passed through as-is to the native endpoint.
"""
config = SnowflakeConfig()
# OpenAI format tool_choice
tool_choice = {"type": "function", "function": {"name": "get_weather"}}
optional_params = {"tool_choice": tool_choice}
transformed_request = config.transform_request(
model="claude-3-5-sonnet",
model="llama3.1-70b",
messages=[{"role": "user", "content": "What's the weather?"}],
optional_params=optional_params,
litellm_params={},
headers={},
)
# Verify tool_choice was transformed to Snowflake format
assert "tool_choice" in transformed_request
assert transformed_request["tool_choice"]["type"] == "tool"
assert transformed_request["tool_choice"]["name"] == [
"get_weather"
] # Array format
assert transformed_request["tool_choice"] == tool_choice
def test_transform_request_with_string_tool_choice(self):
"""
Test that string tool_choice values are transformed to Snowflake object format.
Test that string tool_choice values are passed through as-is to the native endpoint.
Snowflake's API (like Anthropic) requires tool_choice as an object
with a "type" field, not as a bare string. OpenAI's "required" maps
to Snowflake's "any".
The native /chat/completions endpoint accepts OpenAI-style string
tool_choice values directly ("auto", "required", "none").
"""
config = SnowflakeConfig()
expected_mappings = {
"auto": {"type": "auto"},
"required": {"type": "any"},
"none": {"type": "none"},
}
for value, expected in expected_mappings.items():
for value in ["auto", "required", "none"]:
optional_params = {"tool_choice": value}
transformed_request = config.transform_request(
model="claude-3-5-sonnet",
model="llama3.1-70b",
messages=[{"role": "user", "content": "Test"}],
optional_params=optional_params,
litellm_params={},
headers={},
)
assert transformed_request["tool_choice"] == expected, (
f"tool_choice='{value}' should be transformed to {expected}, "
assert transformed_request["tool_choice"] == value, (
f"tool_choice='{value}' should pass through unchanged, "
f"got {transformed_request['tool_choice']}"
)
def test_transform_response_with_tool_calls(self):
"""
Test that Snowflake's content_list with tool_use is transformed to OpenAI format.
Test that standard OpenAI tool_calls response format is parsed correctly.
The native /chat/completions endpoint returns standard OpenAI format.
"""
config = SnowflakeConfig()
# Mock Snowflake response with tool call
mock_snowflake_response = {
mock_response = {
"id": "chatcmpl-123",
"object": "chat.completion",
"model": "llama3.1-70b",
"choices": [
{
"index": 0,
"message": {
"content_list": [
{"type": "text", "text": ""},
"role": "assistant",
"content": None,
"tool_calls": [
{
"type": "tool_use",
"tool_use": {
"tool_use_id": "tooluse_abc123",
"id": "call_abc123",
"type": "function",
"function": {
"name": "get_weather",
"input": {
"location": "Paris, France",
"unit": "celsius",
},
"arguments": json.dumps({"location": "Paris, France", "unit": "celsius"}),
},
},
]
}
}
],
},
"finish_reason": "tool_calls",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
@ -172,7 +155,7 @@ class TestSnowflakeToolTransformation:
response = httpx.Response(
status_code=200,
json=mock_snowflake_response,
json=mock_response,
headers={"Content-Type": "application/json"},
)
@ -183,7 +166,7 @@ class TestSnowflakeToolTransformation:
logging_obj = MagicMock()
result = config.transform_response(
model="claude-3-5-sonnet",
model="llama3.1-70b",
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
@ -194,61 +177,50 @@ class TestSnowflakeToolTransformation:
encoding={},
)
# General assertions
assert isinstance(result, ModelResponse)
assert len(result.choices) == 1
choice = result.choices[0]
assert isinstance(choice, litellm.Choices)
# Message and tool_calls assertions
message = choice.message
assert isinstance(message, litellm.Message)
assert hasattr(message, "tool_calls")
assert isinstance(message.tool_calls, list)
message = result.choices[0].message
assert message.tool_calls is not None
assert len(message.tool_calls) == 1
# Specific tool_call assertions
tool_call = message.tool_calls[0]
assert isinstance(tool_call, litellm.utils.ChatCompletionMessageToolCall)
assert tool_call.id == "tooluse_abc123"
assert tool_call.id == "call_abc123"
assert tool_call.type == "function"
assert tool_call.function.name == "get_weather"
# Verify arguments are properly JSON serialized
arguments = json.loads(tool_call.function.arguments)
assert arguments["location"] == "Paris, France"
assert arguments["unit"] == "celsius"
# Verify content_list was removed and content was set
assert message.content == ""
def test_transform_response_with_mixed_content(self):
"""
Test that responses with both text and tool calls are handled correctly.
Test that responses with both text content and tool calls are parsed correctly.
"""
config = SnowflakeConfig()
# Mock Snowflake response with text and tool call
mock_snowflake_response = {
mock_response = {
"id": "chatcmpl-456",
"object": "chat.completion",
"model": "llama3.1-70b",
"choices": [
{
"index": 0,
"message": {
"content_list": [
"role": "assistant",
"content": "Let me check the weather for you.",
"tool_calls": [
{
"type": "text",
"text": "Let me check the weather for you. ",
},
{
"type": "tool_use",
"tool_use": {
"tool_use_id": "tooluse_xyz789",
"id": "call_xyz789",
"type": "function",
"function": {
"name": "get_weather",
"input": {"location": "Tokyo, Japan"},
"arguments": json.dumps({"location": "Tokyo, Japan"}),
},
},
]
}
}
],
},
"finish_reason": "tool_calls",
}
],
"usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40},
@ -256,7 +228,7 @@ class TestSnowflakeToolTransformation:
response = httpx.Response(
status_code=200,
json=mock_snowflake_response,
json=mock_response,
headers={"Content-Type": "application/json"},
)
@ -267,7 +239,7 @@ class TestSnowflakeToolTransformation:
logging_obj = MagicMock()
result = config.transform_response(
model="claude-3-5-sonnet",
model="llama3.1-70b",
raw_response=response,
model_response=model_response,
logging_obj=logging_obj,
@ -278,11 +250,8 @@ class TestSnowflakeToolTransformation:
encoding={},
)
# Verify text content was extracted
message = result.choices[0].message
assert message.content == "Let me check the weather for you. "
# Verify tool call was also extracted
assert message.content == "Let me check the weather for you."
assert len(message.tool_calls) == 1
assert message.tool_calls[0].function.name == "get_weather"
@ -341,7 +310,7 @@ class TestSnowflakeToolTransformation:
Test that tools and tool_choice are in supported params.
"""
config = SnowflakeConfig()
supported_params = config.get_supported_openai_params("claude-3-5-sonnet")
supported_params = config.get_supported_openai_params("llama3.1-70b")
assert "tools" in supported_params
assert "tool_choice" in supported_params
@ -392,8 +361,8 @@ class TestSnowFlakeCompletion:
assert "00000" in post_kwargs["headers"]["Authorization"]
# account id was used
assert "AAAA-BBBB" in post_kwargs["url"]
# is completion
assert post_kwargs["url"].endswith("cortex/inference:complete")
# uses native endpoint
assert post_kwargs["url"].endswith("cortex/v1/chat/completions")
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
def test_snowflake_pat_key_account_id(self, mock_post):

View file

@ -0,0 +1,718 @@
"""
Tests for Snowflake Cortex native endpoint migration.
Covers:
- SnowflakeConfig with auto-routing:
- Non-Claude models /chat/completions (OpenAI format)
- Claude models /messages (Anthropic format)
Run:
pytest tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py -v
"""
import json
from unittest.mock import MagicMock, patch
import httpx
import pytest
from litellm.llms.snowflake.chat.transformation import (
SnowflakeConfig,
_is_claude_model,
)
from litellm.types.utils import ModelResponse
# ─── Fixtures ──────────────────────────────────────────────────────────────
ACCOUNT_ID = "myaccount"
API_BASE = f"https://{ACCOUNT_ID}.snowflakecomputing.com"
PAT_TOKEN = "pat/my-secret-pat-token"
JWT_TOKEN = "eyJhbGciOiJSUzI1NiJ9.test"
def _mock_logging():
m = MagicMock()
m.post_call = MagicMock()
return m
def _make_openai_response(content: str = "Hello!") -> httpx.Response:
body = {
"id": "chatcmpl-abc123",
"object": "chat.completion",
"model": "llama3.1-70b",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
return httpx.Response(200, json=body)
def _make_anthropic_response(content: str = "Hello!") -> httpx.Response:
body = {
"id": "msg_abc123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": content}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
}
return httpx.Response(200, json=body)
# ─── SnowflakeConfig (OpenAI-compatible) ───────────────────────────────────
class TestSnowflakeConfigURL:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_url_with_account_id_in_optional_params(self):
optional_params = {"account_id": ACCOUNT_ID}
url = self.cfg.get_complete_url(
api_base=None,
api_key=JWT_TOKEN,
model="snowflake/llama3.1-70b",
optional_params=optional_params,
litellm_params={},
)
assert url == f"https://{ACCOUNT_ID}.snowflakecomputing.com/api/v2/cortex/v1/chat/completions"
def test_url_with_explicit_api_base(self):
url = self.cfg.get_complete_url(
api_base=API_BASE,
api_key=JWT_TOKEN,
model="snowflake/llama3.1-70b",
optional_params={},
litellm_params={},
)
assert url.endswith("/api/v2/cortex/v1/chat/completions")
assert "cortex/inference:complete" not in url
def test_url_never_uses_legacy_endpoint(self):
url = self.cfg.get_complete_url(
api_base=API_BASE,
api_key=JWT_TOKEN,
model="snowflake/llama3.1-70b",
optional_params={},
litellm_params={},
)
assert "inference:complete" not in url
assert "/v1/chat/completions" in url
def test_url_works_for_claude_models(self):
url = self.cfg.get_complete_url(
api_base=API_BASE,
api_key=JWT_TOKEN,
model="snowflake/claude-sonnet-4-5",
optional_params={},
litellm_params={},
)
assert "/cortex/v1/messages" in url
def test_url_works_for_llama_models(self):
url = self.cfg.get_complete_url(
api_base=API_BASE,
api_key=JWT_TOKEN,
model="snowflake/llama3.1-70b",
optional_params={},
litellm_params={},
)
assert "/cortex/v1/chat/completions" in url
class TestSnowflakeConfigAuth:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_pat_auth_strips_prefix_and_sets_header(self):
headers = self.cfg.validate_environment(
headers={},
model="snowflake/llama3.1-70b",
messages=[],
optional_params={},
litellm_params={},
api_key=PAT_TOKEN,
)
assert headers["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN"
assert headers["Authorization"] == "Bearer my-secret-pat-token"
def test_jwt_auth_sets_keypair_header(self):
headers = self.cfg.validate_environment(
headers={},
model="snowflake/llama3.1-70b",
messages=[],
optional_params={},
litellm_params={},
api_key=JWT_TOKEN,
)
assert headers["X-Snowflake-Authorization-Token-Type"] == "KEYPAIR_JWT"
assert headers["Authorization"] == f"Bearer {JWT_TOKEN}"
def test_missing_api_key_raises(self):
with pytest.raises(ValueError, match="Missing Snowflake JWT key"):
self.cfg.validate_environment(
headers={},
model="snowflake/llama3.1-70b",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)
class TestSnowflakeConfigRequest:
def setup_method(self):
self.cfg = SnowflakeConfig()
self.messages = [{"role": "user", "content": "hello"}]
def test_request_uses_openai_tool_format(self):
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
},
}
]
body = self.cfg.transform_request(
model="snowflake/llama3.1-70b",
messages=self.messages,
optional_params={"tools": tools},
litellm_params={},
headers={},
)
assert body["tools"] == tools
assert "tool_spec" not in json.dumps(body)
def test_stream_defaults_to_false(self):
body = self.cfg.transform_request(
model="snowflake/llama3.1-70b",
messages=self.messages,
optional_params={},
litellm_params={},
headers={},
)
assert body["stream"] is False
def test_stream_true_passes_through(self):
body = self.cfg.transform_request(
model="snowflake/llama3.1-70b",
messages=self.messages,
optional_params={"stream": True},
litellm_params={},
headers={},
)
assert body["stream"] is True
def test_supported_params_includes_stream(self):
params = self.cfg.get_supported_openai_params("snowflake/llama3.1-70b")
assert "stream" in params
def test_no_content_list_in_request(self):
body = self.cfg.transform_request(
model="snowflake/llama3.1-70b",
messages=self.messages,
optional_params={},
litellm_params={},
headers={},
)
assert "content_list" not in body
class TestSnowflakeConfigResponse:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_standard_response_parsed(self):
raw = _make_openai_response("Hello from Snowflake!")
result = self.cfg.transform_response(
model="snowflake/llama3.1-70b",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[{"role": "user", "content": "hi"}],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.choices[0].message.content == "Hello from Snowflake!"
assert result.model.startswith("snowflake/")
def test_model_prefixed_with_snowflake(self):
raw = _make_openai_response()
result = self.cfg.transform_response(
model="snowflake/llama3.1-70b",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.model.startswith("snowflake/")
# ─── SnowflakeConfig ────────────────────────────────────────
class TestAnthropicConfigURL:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_url_routes_to_messages_endpoint(self):
url = self.cfg.get_complete_url(
api_base=API_BASE,
api_key=PAT_TOKEN,
model="snowflake/claude-sonnet-4-5",
optional_params={},
litellm_params={},
)
assert url.endswith("/api/v2/cortex/v1/messages")
assert "chat/completions" not in url
assert "inference:complete" not in url
def test_url_with_account_id(self):
url = self.cfg.get_complete_url(
api_base=None,
api_key=PAT_TOKEN,
model="snowflake/claude-sonnet-4-5",
optional_params={"account_id": ACCOUNT_ID},
litellm_params={},
)
assert f"https://{ACCOUNT_ID}.snowflakecomputing.com/api/v2/cortex/v1/messages" == url
class TestAnthropicConfigAuth:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_anthropic_version_header_set(self):
headers = self.cfg.validate_environment(
headers={},
model="snowflake/claude-sonnet-4-5",
messages=[],
optional_params={},
litellm_params={},
api_key=PAT_TOKEN,
)
assert headers["anthropic-version"] == "2023-06-01"
def test_pat_auth_and_anthropic_version_combined(self):
headers = self.cfg.validate_environment(
headers={},
model="snowflake/claude-sonnet-4-5",
messages=[],
optional_params={},
litellm_params={},
api_key=PAT_TOKEN,
)
assert headers["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN"
assert headers["anthropic-version"] == "2023-06-01"
assert "Bearer" in headers["Authorization"]
class TestAnthropicConfigRequest:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_system_message_extracted_to_top_level(self):
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
assert body["system"] == "You are helpful."
assert all(m["role"] != "system" for m in body["messages"])
assert body["messages"][0] == {"role": "user", "content": "Hello"}
def test_model_prefix_stripped(self):
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
optional_params={},
litellm_params={},
headers={},
)
assert body["model"] == "claude-sonnet-4-5"
assert "snowflake/" not in body["model"]
def test_max_tokens_defaulted_when_missing(self):
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
optional_params={},
litellm_params={},
headers={},
)
assert "max_tokens" in body
assert body["max_tokens"] == 4096
def test_max_tokens_not_overridden_when_provided(self):
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
optional_params={"max_tokens": 500},
litellm_params={},
headers={},
)
assert body["max_tokens"] == 500
def test_no_system_key_when_no_system_message(self):
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
optional_params={},
litellm_params={},
headers={},
)
assert "system" not in body
class TestAnthropicConfigResponse:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_anthropic_response_to_openai_format(self):
raw = _make_anthropic_response("Hi there!")
result = self.cfg.transform_response(
model="snowflake/claude-sonnet-4-5",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[{"role": "user", "content": "hi"}],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.choices[0].message.content == "Hi there!"
assert result.choices[0].finish_reason == "stop"
def test_usage_tokens_mapped(self):
raw = _make_anthropic_response()
result = self.cfg.transform_response(
model="snowflake/claude-sonnet-4-5",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.usage.prompt_tokens == 10
assert result.usage.completion_tokens == 5
assert result.usage.total_tokens == 15
def test_stop_reason_end_turn_maps_to_stop(self):
raw = _make_anthropic_response()
result = self.cfg.transform_response(
model="snowflake/claude-sonnet-4-5",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.choices[0].finish_reason == "stop"
def test_tool_use_block_mapped_to_tool_calls(self):
body = {
"id": "msg_tool",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [
{
"type": "tool_use",
"id": "toolu_01",
"name": "get_weather",
"input": {"city": "Paris"},
}
],
"stop_reason": "tool_use",
"usage": {"input_tokens": 20, "output_tokens": 10},
}
raw = httpx.Response(200, json=body)
result = self.cfg.transform_response(
model="snowflake/claude-sonnet-4-5",
raw_response=raw,
model_response=ModelResponse(),
logging_obj=_mock_logging(),
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
assert result.choices[0].finish_reason == "tool_calls"
tool_calls = result.choices[0].message.tool_calls
assert len(tool_calls) == 1
assert tool_calls[0].function.name == "get_weather"
assert json.loads(tool_calls[0].function.arguments) == {"city": "Paris"}
# ─── Model detection helper ────────────────────────────────────────────────
class TestIsClaudeModel:
def test_claude_model_detected(self):
assert _is_claude_model("snowflake/claude-sonnet-4-5") is True
assert _is_claude_model("claude-3-haiku") is True
assert _is_claude_model("snowflake/claude-opus-4") is True
def test_non_claude_not_detected(self):
assert _is_claude_model("snowflake/llama3.1-70b") is False
assert _is_claude_model("snowflake/mistral-large") is False
assert _is_claude_model("snowflake/deepseek-r1") is False
assert _is_claude_model("snowflake/snowflake-arctic") is False
# ─── Anthropic Tool Transformation Tests ──────────────────────────────────
class TestAnthropicToolTransformation:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_openai_tools_converted_to_anthropic_format(self):
messages = [{"role": "user", "content": "What's the weather?"}]
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get current weather",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
},
}
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={"tools": tools},
litellm_params={},
headers={},
)
assert len(body["tools"]) == 1
tool = body["tools"][0]
assert tool["name"] == "get_weather"
assert tool["description"] == "Get current weather"
assert "input_schema" in tool
assert tool["input_schema"]["properties"]["city"]["type"] == "string"
assert "function" not in tool
assert "type" not in tool
def test_tools_already_in_anthropic_format_pass_through(self):
messages = [{"role": "user", "content": "hi"}]
tools = [{"name": "my_tool", "input_schema": {"type": "object", "properties": {}}}]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={"tools": tools},
litellm_params={},
headers={},
)
assert body["tools"] == tools
class TestAnthropicMultiTurnToolMessages:
def setup_method(self):
self.cfg = SnowflakeConfig()
def test_assistant_tool_calls_converted_to_tool_use_blocks(self):
messages = [
{"role": "user", "content": "What's the weather in Paris?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"city": "Paris"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "call_123",
"content": "Sunny, 22°C",
},
{"role": "user", "content": "Thanks!"},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
msgs = body["messages"]
assert msgs[0] == {"role": "user", "content": "What's the weather in Paris?"}
assistant_msg = msgs[1]
assert assistant_msg["role"] == "assistant"
assert isinstance(assistant_msg["content"], list)
assert assistant_msg["content"][0]["type"] == "tool_use"
assert assistant_msg["content"][0]["id"] == "call_123"
assert assistant_msg["content"][0]["name"] == "get_weather"
assert assistant_msg["content"][0]["input"] == {"city": "Paris"}
tool_result_msg = msgs[2]
assert tool_result_msg["role"] == "user"
assert tool_result_msg["content"][0]["type"] == "tool_result"
assert tool_result_msg["content"][0]["tool_use_id"] == "call_123"
assert tool_result_msg["content"][0]["content"] == "Sunny, 22°C"
assert msgs[3] == {"role": "user", "content": "Thanks!"}
def test_assistant_with_text_and_tool_calls(self):
messages = [
{"role": "user", "content": "Check weather"},
{
"role": "assistant",
"content": "Let me check that for you.",
"tool_calls": [
{
"id": "call_456",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"city": "London"}',
},
}
],
},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
assistant_msg = body["messages"][1]
assert assistant_msg["content"][0] == {"type": "text", "text": "Let me check that for you."}
assert assistant_msg["content"][1]["type"] == "tool_use"
assert assistant_msg["content"][1]["name"] == "get_weather"
def test_tool_role_never_in_output(self):
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}],
},
{"role": "tool", "tool_call_id": "c1", "content": "result"},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
for msg in body["messages"]:
assert msg["role"] != "tool"
def test_malformed_json_in_tool_arguments_handled_gracefully(self):
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_bad",
"type": "function",
"function": {"name": "broken_tool", "arguments": "not valid json{{{"},
}
],
},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
assistant_msg = body["messages"][1]
tool_use_block = assistant_msg["content"][0]
assert tool_use_block["type"] == "tool_use"
assert tool_use_block["name"] == "broken_tool"
assert tool_use_block["input"] == {}
def test_non_string_tool_arguments_pass_through(self):
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_dict",
"type": "function",
"function": {"name": "dict_tool", "arguments": {"already": "parsed"}},
}
],
},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
tool_use_block = body["messages"][1]["content"][0]
assert tool_use_block["input"] == {"already": "parsed"}
def test_tool_result_with_non_string_content(self):
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}],
},
{"role": "tool", "tool_call_id": "c1", "content": {"result_key": "result_value"}},
]
body = self.cfg.transform_request(
model="snowflake/claude-sonnet-4-5",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
tool_result = body["messages"][2]["content"][0]
assert tool_result["type"] == "tool_result"
assert json.loads(tool_result["content"]) == {"result_key": "result_value"}

View file

@ -0,0 +1,306 @@
import json
from unittest.mock import MagicMock
import pytest
class TestVoyageMultimodalEmbeddings:
def test_multimodal_model_detection(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(
"voyage-multimodal-3.5"
)
assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(
"voyage-multimodal-3"
)
assert not VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings("voyage-4")
def test_multimodal_embedding_url_generation(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
assert (
config.get_complete_url(None, None, "voyage-multimodal-3.5", {}, {})
== "https://api.voyageai.com/v1/multimodalembeddings"
)
assert (
config.get_complete_url(
"https://custom.api.com", None, "voyage-multimodal-3.5", {}, {}
)
== "https://custom.api.com/multimodalembeddings"
)
assert (
config.get_complete_url(
"https://custom.api.com/multimodalembeddings",
None,
"voyage-multimodal-3.5",
{},
{},
)
== "https://custom.api.com/multimodalembeddings"
)
def test_multimodal_embedding_request_transformation(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
data_uri = "data:image/png;base64,AAAA"
request = config.transform_embedding_request(
"voyage-multimodal-3.5",
[
{
"content": [
{"type": "text", "text": "Describe this"},
{"type": "image_url", "image_url": {"url": data_uri}},
{"type": "image_url", "image_url": "https://example.com/a.png"},
]
}
],
{"input_type": "document", "output_dimension": 512},
{},
)
assert request["model"] == "voyage-multimodal-3.5"
assert "inputs" in request
assert "input" not in request
assert request["input_type"] == "document"
assert request["output_dimension"] == 512
assert request["inputs"][0]["content"][1] == {
"type": "image_base64",
"image_base64": "AAAA",
}
assert request["inputs"][0]["content"][2] == {
"type": "image_url",
"image_url": "https://example.com/a.png",
}
def test_multimodal_embedding_string_input_transformation(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
request = config.transform_embedding_request(
"voyage-multimodal-3.5", "hello", {}, {}
)
assert request["inputs"] == [
{"content": [{"type": "text", "text": "hello"}]}
]
def test_multimodal_embedding_response_transformation(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
from litellm.types.utils import EmbeddingResponse
config = VoyageMultimodalEmbeddingConfig()
response_payload = {
"object": "list",
"data": [
{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}
],
"model": "voyage-multimodal-3.5",
"usage": {
"text_tokens": 2,
"image_pixels": 0,
"video_pixels": 0,
"total_tokens": 2,
},
}
raw_response = MagicMock()
raw_response.json.return_value = response_payload
raw_response.status_code = 200
raw_response.text = json.dumps(response_payload)
model_response = EmbeddingResponse()
transformed = config.transform_embedding_response(
"voyage-multimodal-3.5", raw_response, model_response, MagicMock()
)
assert transformed.model == "voyage-multimodal-3.5"
assert transformed.object == "list"
assert transformed.data == response_payload["data"]
assert transformed.usage.prompt_tokens == 2
assert transformed.usage.total_tokens == 2
def test_provider_config_manager_routes_multimodal_models(self):
import litellm
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
from litellm.utils import ProviderConfigManager
config = ProviderConfigManager.get_provider_embedding_config(
model="voyage-multimodal-3.5", provider=litellm.LlmProviders.VOYAGE
)
assert isinstance(config, VoyageMultimodalEmbeddingConfig)
def test_map_openai_params_dimensions(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
assert config.get_supported_openai_params("voyage-multimodal-3.5") == [
"dimensions"
]
optional_params = config.map_openai_params(
{"dimensions": 512}, {}, "voyage-multimodal-3.5", False
)
assert optional_params == {"output_dimension": 512}
assert (
config.map_openai_params({}, {}, "voyage-multimodal-3.5", False) == {}
)
def test_validate_environment_uses_api_key(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
headers = config.validate_environment(
{}, "voyage-multimodal-3.5", [], {}, {}, api_key="test-key"
)
assert headers == {"Authorization": "Bearer test-key"}
def test_validate_environment_uses_secret_fallback(self, monkeypatch):
import litellm.llms.voyage.embedding.transformation_multimodal as module
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
def fake_get_secret(name):
return "secret-key" if name == "VOYAGE_AI_API_KEY" else None
monkeypatch.setattr(module, "get_secret_str", fake_get_secret)
config = VoyageMultimodalEmbeddingConfig()
headers = config.validate_environment(
{}, "voyage-multimodal-3.5", [], {}, {}, api_key=None
)
assert headers == {"Authorization": "Bearer secret-key"}
def test_validate_environment_raises_without_api_key(self, monkeypatch):
import litellm.llms.voyage.embedding.transformation_multimodal as module
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
monkeypatch.setattr(module, "get_secret_str", lambda name: None)
config = VoyageMultimodalEmbeddingConfig()
with pytest.raises(ValueError) as exc_info:
config.validate_environment(
{}, "voyage-multimodal-3.5", [], {}, {}, api_key=None
)
assert "VOYAGE_API_KEY" in str(exc_info.value)
def test_normalize_image_url_dict_missing_url_raises(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
with pytest.raises(ValueError) as exc_info:
config._normalize_content_item({"type": "image_url", "image_url": {}})
assert "image_url" in str(exc_info.value)
def test_is_multimodal_embeddings_helper(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(
"voyage-multimodal-3"
)
assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(
"VOYAGE-MULTIMODAL-3.5"
)
assert not VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(
"voyage-3.5"
)
def test_utils_routing_via_provider_config_and_dimensions(self):
import litellm
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
from litellm.utils import (
ProviderConfigManager,
get_optional_params_embeddings,
)
config = ProviderConfigManager.get_provider_embedding_config(
model="voyage-multimodal-3.5", provider=litellm.LlmProviders.VOYAGE
)
assert isinstance(config, VoyageMultimodalEmbeddingConfig)
optional_params = get_optional_params_embeddings(
model="voyage-multimodal-3.5",
dimensions=1024,
custom_llm_provider="voyage",
drop_params=True,
)
assert optional_params.get("output_dimension") == 1024
def test_get_supported_openai_params_voyage_routes_multimodal(self):
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
)
multimodal_params = get_supported_openai_params(
model="voyage-multimodal-3.5",
custom_llm_provider="voyage",
request_type="embeddings",
)
assert multimodal_params == ["dimensions"]
standard_params = get_supported_openai_params(
model="voyage-3.5",
custom_llm_provider="voyage",
request_type="embeddings",
)
assert "dimensions" in standard_params
assert "encoding_format" in standard_params
def test_passthrough_non_content_input(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
)
config = VoyageMultimodalEmbeddingConfig()
request = config.transform_embedding_request(
"voyage-multimodal-3.5", [{"foo": "bar"}], {}, {}
)
assert request["inputs"] == [{"foo": "bar"}]
def test_error_response_transformation_and_error_class(self):
from litellm.llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig,
VoyageMultimodalEmbeddingError,
)
from litellm.types.utils import EmbeddingResponse
config = VoyageMultimodalEmbeddingConfig()
raw_response = MagicMock()
raw_response.json.side_effect = ValueError("not json")
raw_response.status_code = 400
raw_response.text = "bad request"
with pytest.raises(VoyageMultimodalEmbeddingError) as exc_info:
config.transform_embedding_response(
"voyage-multimodal-3.5", raw_response, EmbeddingResponse(), MagicMock()
)
assert exc_info.value.status_code == 400
assert exc_info.value.message == "bad request"
error = config.get_error_class("rate limited", 429, {"x-test": "1"})
assert isinstance(error, VoyageMultimodalEmbeddingError)
assert error.status_code == 429
assert error.message == "rate limited"

View file

@ -1221,6 +1221,138 @@ async def test_health_endpoint_filters_model_list_by_user_access():
}, f"health_endpoint did not scope model_list to caller access: {returned_names}"
@pytest.mark.asyncio
async def test_health_endpoint_keeps_full_model_list_for_all_proxy_models():
"""
A key granted all model permissions carries the literal
"all-proxy-models" entry in user_api_key_dict.models. It matches no real
model_name, so the access filter must be skipped entirely; otherwise the
model list filters down to nothing and /health reports 0/0 counts.
"""
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-a"},
},
{
"model_name": "model-b",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-b"},
},
]
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
models=[SpecialModelNames.all_proxy_models.value],
)
captured: dict = {}
async def fake_perform(**kwargs):
captured["model_list"] = kwargs["model_list"]
return {
"healthy_endpoints": [],
"unhealthy_endpoints": [],
"healthy_count": 0,
"unhealthy_count": 0,
}
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", {}),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
patch(
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
side_effect=fake_perform,
),
):
from fastapi import Response
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
returned_names = {m["model_name"] for m in captured["model_list"]}
assert returned_names == {
"model-a",
"model-b",
}, f"all-proxy-models key should health-check every model: {returned_names}"
@pytest.mark.asyncio
async def test_health_endpoint_resolves_all_team_models_to_team_allowlist():
"""
A key granted "all-team-models" carries the literal sentinel in
user_api_key_dict.models, which matches no real model_name. With a
team_id the sentinel must resolve to the team's allowlist (same
semantics as get_key_models); otherwise the filter would zero out the
model list just like the all-proxy-models case.
"""
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-a"},
},
{
"model_name": "model-b",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-b"},
},
]
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
models=[SpecialModelNames.all_team_models.value],
team_id="team-1",
team_models=["model-b"],
)
captured: dict = {}
async def fake_perform(**kwargs):
captured["model_list"] = kwargs["model_list"]
return {
"healthy_endpoints": [],
"unhealthy_endpoints": [],
"healthy_count": 0,
"unhealthy_count": 0,
}
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", {}),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
patch(
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
side_effect=fake_perform,
),
):
from fastapi import Response
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
returned_names = {m["model_name"] for m in captured["model_list"]}
assert returned_names == {
"model-b"
}, f"all-team-models key should health-check the team's models: {returned_names}"
@pytest.mark.asyncio
async def test_health_endpoint_filters_background_cache_by_user_access():
"""

View file

@ -6555,14 +6555,20 @@ async def test_reset_key_spend_success(monkeypatch):
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache,
patch(
"litellm.proxy.proxy_server._invalidate_spend_counter"
) as mock_invalidate,
):
mock_hash_token.return_value = hashed_key
mock_check_admin.return_value = None
mock_delete_cache.return_value = None
# Mock spend_counter_cache to verify direct cache set instead of
# _invalidate_spend_counter (removed in favour of atomic cache write).
mock_spend_counter_cache = MagicMock()
mock_spend_counter_cache.redis_cache = None
monkeypatch.setattr(
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
@ -6582,7 +6588,9 @@ async def test_reset_key_spend_success(monkeypatch):
assert response["max_budget"] == 200.0
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
mock_delete_cache.assert_awaited_once()
mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}")
mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(
key=f"spend:key:{hashed_key}", value=50.0, ttl=60
)
@pytest.mark.asyncio
@ -11853,83 +11861,83 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption(
assert str(code) == "400"
assert "cannot exceed" in msg.lower()
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
"""
When budget_duration is explicitly set to null, prepare_key_update_data
should produce budget_duration=None and budget_reset_at=None so Prisma
clears them in the DB.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration=None)
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" in result
assert result["budget_duration"] is None
assert "budget_reset_at" in result
assert result["budget_reset_at"] is None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
"""
When budget_duration is NOT sent in the request (unset), it should not
appear in the result dict at all the existing DB value stays unchanged.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" not in result
assert "budget_reset_at" not in result
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
"""
When budget_duration is set to a valid duration string, both
budget_duration and budget_reset_at should be populated.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert result["budget_duration"] == "30d"
assert result["budget_reset_at"] is not None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
"""
When budget_duration is explicitly set to null, prepare_key_update_data
should produce budget_duration=None and budget_reset_at=None so Prisma
clears them in the DB.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration=None)
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" in result
assert result["budget_duration"] is None
assert "budget_reset_at" in result
assert result["budget_reset_at"] is None
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
"""
When budget_duration is NOT sent in the request (unset), it should not
appear in the result dict at all the existing DB value stays unchanged.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert "budget_duration" not in result
assert "budget_reset_at" not in result
@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
"""
When budget_duration is set to a valid duration string, both
budget_duration and budget_reset_at should be populated.
"""
existing_key = LiteLLM_VerificationToken(
token="test-token",
key_alias="test-key",
models=[],
user_id="test-user",
team_id=None,
metadata={},
)
update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")
result = await prepare_key_update_data(
data=update_request, existing_key_row=existing_key
)
assert result["budget_duration"] == "30d"
assert result["budget_reset_at"] is not None

View file

@ -1609,7 +1609,8 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name):
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging,
patch(
"litellm.proxy.management_endpoints.team_endpoints._cache_team_object"
"litellm.proxy.management_endpoints.team_endpoints._cache_team_object",
new_callable=AsyncMock,
) as mock_cache_team,
):
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
@ -1618,7 +1619,7 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name):
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(
return_value=updated_team
)
mock_cache_team.return_value = None
mock_prisma_client.db.execute_raw = AsyncMock(return_value=None)
if endpoint_name == "team_model_add":
await team_model_add(

View file

@ -0,0 +1,83 @@
"""
Tests for atomic team model operations during BYOK model creation.
Regression tests for https://github.com/BerriAI/litellm/issues/22594
Concurrent BYOK model creates must not overwrite each other's entries
in team.models.
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.proxy._types import (
LitellmUserRoles,
TeamModelAddRequest,
UserAPIKeyAuth,
)
class TestTeamModelAddAtomicAppend:
"""Verify team_model_add uses atomic SQL for the models array append."""
@pytest.mark.asyncio
async def test_uses_atomic_array_append_with_dedup(self):
"""team_model_add must call execute_raw with DISTINCT unnest SQL."""
from unittest.mock import patch
from litellm.proxy.management_endpoints.team_endpoints import team_model_add
mock_request = MagicMock()
mock_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user"
)
existing_team = MagicMock()
existing_team.model_dump.return_value = {
"team_id": "team-1",
"models": ["existing-model"],
}
updated_team = MagicMock()
updated_team.team_id = "team-1"
updated_team.model_dump.return_value = {
"team_id": "team-1",
"models": ["existing-model", "new-model"],
}
with (
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch(
"litellm.proxy.management_endpoints.team_endpoints._cache_team_object",
new_callable=AsyncMock,
),
patch("litellm.proxy.proxy_server.user_api_key_cache"),
patch("litellm.proxy.proxy_server.proxy_logging_obj"),
):
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
return_value=existing_team
)
mock_prisma.db.execute_raw = AsyncMock(return_value=None)
mock_prisma.db.litellm_teamtable.update = AsyncMock(
return_value=updated_team
)
await team_model_add(
data=TeamModelAddRequest(team_id="team-1", models=["new-model"]),
http_request=mock_request,
user_api_key_dict=mock_user,
)
mock_prisma.db.execute_raw.assert_called_once()
sql = mock_prisma.db.execute_raw.call_args[0][0]
assert "DISTINCT unnest" in sql
assert "all-proxy-models" in sql
assert mock_prisma.db.execute_raw.call_args[0][1] == ["new-model"]
assert mock_prisma.db.execute_raw.call_args[0][2] == "team-1"
# Should use write-routed update to re-fetch, not find_unique
mock_prisma.db.litellm_teamtable.update.assert_called_once()

View file

@ -9,7 +9,6 @@ Pins covered:
- ``initialize``
- ``load_from_azure_key_vault``
- ``cost_tracking``
- ``check_request_disconnection``
- ``_resolve_typed_dict_type``
- ``_resolve_pydantic_type``
- ``get_litellm_model_info``
@ -26,7 +25,7 @@ from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI, HTTPException
from fastapi import FastAPI
from pydantic import BaseModel
from typing_extensions import TypedDict
@ -35,7 +34,6 @@ from litellm.proxy.proxy_server import (
_initialize_shared_aiohttp_session,
_resolve_pydantic_type,
_resolve_typed_dict_type,
check_request_disconnection,
cleanup_router_config_variables,
cost_tracking,
get_litellm_model_info,
@ -324,62 +322,6 @@ def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):
assert litellm._async_success_callback == []
# ---------------------------------------------------------------------------
# check_request_disconnection
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_check_request_disconnection_cancels_task_and_raises_499(monkeypatch):
monkeypatch.setattr(ps.asyncio, "sleep", AsyncMock(return_value=None))
request = MagicMock()
request.is_disconnected = AsyncMock(return_value=True)
task = MagicMock()
raised_status = None
try:
await check_request_disconnection(request=request, llm_api_call_task=task)
except HTTPException as exc:
raised_status = exc.status_code
observed = {
"raised_status": raised_status,
"cancel_called": task.cancel.called,
"is_async": inspect.iscoroutinefunction(check_request_disconnection),
}
assert normalize(observed) == {
"raised_status": 499,
"cancel_called": True,
"is_async": True,
}
@pytest.mark.asyncio
async def test_check_request_disconnection_invalid_when_connected_times_out(monkeypatch):
"""With a connected request the function loops for up to 10 minutes —
wrap in wait_for and assert it times out. Patch ``asyncio.sleep`` so the
loop spins without real wall-clock waits."""
import litellm.proxy.proxy_server as ps
request = MagicMock()
request.is_disconnected = AsyncMock(return_value=False)
task = MagicMock()
_real_sleep = asyncio.sleep
async def _instant_sleep(_seconds):
await _real_sleep(0)
monkeypatch.setattr(ps.asyncio, "sleep", _instant_sleep)
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(
check_request_disconnection(request=request, llm_api_call_task=task),
timeout=0.05,
)
# ---------------------------------------------------------------------------
# _resolve_typed_dict_type
# ---------------------------------------------------------------------------

View file

@ -980,7 +980,7 @@ async def test_ProxyConfig__update_llm_router_bad_proxy_logging_raises(monkeypat
# Passing None for proxy_logging_obj triggers AttributeError in _add_general_settings_from_db_config
# when it calls proxy_logging_obj.update_values.
with pytest.raises(AttributeError):
await pc._update_llm_router(new_models=None, proxy_logging_obj=None) # type: ignore[arg-type]
await pc._update_llm_router(new_models=[], proxy_logging_obj=None) # type: ignore[arg-type]
# ---------------------------------------------------------------------------

View file

@ -1,6 +1,7 @@
import asyncio
import copy
import datetime
from typing import AsyncGenerator
from typing import AsyncGenerator, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -15,6 +16,8 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
ProxyConfig,
_await_llm_call_cancelling_on_disconnect,
_cancel_llm_call_on_client_disconnect,
_extract_error_from_sse_chunk,
_get_cost_breakdown_from_logging_obj,
_has_attribute_error_in_chain,
@ -2412,6 +2415,77 @@ class TestHandleLLMApiExceptionDictDetail:
assert proxy_exc.code == "500"
class TestHandleLLMApiExceptionRetryAfter:
"""RouterRateLimitError cooldown_time must surface as a retry-after header."""
async def _invoke(self, exc: Exception, callback_headers: Optional[dict] = None):
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
processor = ProxyBaseLLMRequestProcessing(data={})
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
return_value=callback_headers or {}
)
try:
await processor._handle_llm_api_exception(
e=exc,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
except ProxyException as raised:
return raised
raise AssertionError("ProxyException was not raised")
async def test_handle_llm_api_exception_sets_retry_after_from_cooldown_time(self):
from litellm.types.router import RouterRateLimitError
exc = RouterRateLimitError(
model="gpt-4",
cooldown_time=42.3,
enable_pre_call_checks=False,
cooldown_list=[],
)
proxy_exc = await self._invoke(exc)
assert proxy_exc.headers["retry-after"] == "43"
assert proxy_exc.code == "429"
async def test_handle_llm_api_exception_skips_retry_after_when_cooldown_is_zero(
self,
):
from litellm.types.router import RouterRateLimitError
exc = RouterRateLimitError(
model="gpt-4",
cooldown_time=0,
enable_pre_call_checks=False,
cooldown_list=[],
)
proxy_exc = await self._invoke(exc)
assert "retry-after" not in proxy_exc.headers
async def test_handle_llm_api_exception_no_retry_after_for_plain_exception(self):
proxy_exc = await self._invoke(ValueError("some other failure"))
assert "retry-after" not in proxy_exc.headers
async def test_handle_llm_api_exception_retry_after_survives_callback_headers(self):
from litellm.types.router import RouterRateLimitError
exc = RouterRateLimitError(
model="gpt-4",
cooldown_time=42.3,
enable_pre_call_checks=False,
cooldown_list=[],
)
proxy_exc = await self._invoke(
exc, callback_headers={"retry-after": "", "x-custom": "1"}
)
assert proxy_exc.headers["retry-after"] == "43"
assert proxy_exc.headers["x-custom"] == "1"
class TestAsyncStreamingDataGeneratorFastPath:
"""Fast/slow path branching in async_streaming_data_generator."""
@ -2482,6 +2556,197 @@ class TestAsyncStreamingDataGeneratorFastPath:
ProxyLogging._callback_capabilities_cache.clear()
class TestCancelOnDisconnect:
"""
Coverage for the opt-in `general_settings.cancel_on_disconnect` flag:
cancelling the in-flight upstream LLM call when the HTTP client disconnects
(issue #13774), without changing the default code path and without skipping
failure accounting (post_call_failure_hook) on the resulting 499.
"""
def _request(self, messages: list) -> Request:
async def receive():
if messages:
return messages.pop(0)
await asyncio.Event().wait()
return Request(scope={"type": "http", "headers": []}, receive=receive)
async def test_monitor_cancels_llm_call_and_sets_event_on_disconnect(self):
request = self._request(
[
{"type": "http.request", "body": b"", "more_body": False},
{"type": "http.disconnect"},
]
)
llm_call = asyncio.get_running_loop().create_future()
disconnect_event = asyncio.Event()
await _cancel_llm_call_on_client_disconnect(
request, llm_call, disconnect_event
)
assert llm_call.cancelled()
assert disconnect_event.is_set()
async def test_monitor_is_noop_while_client_stays_connected(self):
request = self._request(
[{"type": "http.request", "body": b"", "more_body": False}]
)
llm_call = asyncio.get_running_loop().create_future()
disconnect_event = asyncio.Event()
monitor = asyncio.create_task(
_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
)
await asyncio.sleep(0.01)
assert not monitor.done()
assert not llm_call.cancelled()
assert not disconnect_event.is_set()
monitor.cancel()
async def test_monitor_survives_receive_failure_without_cancelling(self):
"""If request.receive() fails (e.g. transport reset) the watcher must
degrade to a no-op instead of crashing or cancelling the LLM call."""
async def receive():
raise RuntimeError("transport reset")
request = Request(scope={"type": "http", "headers": []}, receive=receive)
llm_call = asyncio.get_running_loop().create_future()
disconnect_event = asyncio.Event()
await _cancel_llm_call_on_client_disconnect(
request, llm_call, disconnect_event
)
assert not llm_call.cancelled()
assert not disconnect_event.is_set()
async def test_cancellation_without_disconnect_reraises_cancelled_error(self):
"""A CancelledError that is NOT client-initiated (e.g. server shutdown)
must propagate as-is instead of being masked as a 499."""
request = self._request([])
llm_call = asyncio.get_running_loop().create_future()
llm_call.cancel()
with pytest.raises(asyncio.CancelledError):
await _await_llm_call_cancelling_on_disconnect(request, llm_call)
async def _drive_base_process_llm_request(
self, monkeypatch, general_settings: dict, llm_call, request: Request
):
from litellm.proxy._types import UserAPIKeyAuth
logging_obj = MagicMock()
logging_obj.litellm_call_id = "test-cancel-on-disconnect"
logging_obj._defer_async_logging = False
logging_obj._on_deferred_stream_complete = None
logging_obj.cost_breakdown = None
processor = ProxyBaseLLMRequestProcessing(
data={"model": "fake-model", "litellm_logging_obj": logging_obj}
)
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
proxy_logging_obj.post_call_success_hook = AsyncMock(
side_effect=lambda data, user_api_key_dict, response: response
)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
return_value=None
)
async def fake_route_request(**kwargs):
return llm_call()
monkeypatch.setattr(
litellm.proxy.common_request_processing,
"route_request",
fake_route_request,
)
return await processor.base_process_llm_request(
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
route_type="acompletion",
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
proxy_config=MagicMock(spec=ProxyConfig),
skip_pre_call_logic=True,
)
async def test_disconnect_ignored_when_flag_disabled(self, monkeypatch):
upstream_cancelled = asyncio.Event()
model_response = litellm.ModelResponse()
async def llm_call():
try:
await asyncio.sleep(0.05)
return model_response
except asyncio.CancelledError:
upstream_cancelled.set()
raise
result = await self._drive_base_process_llm_request(
monkeypatch,
general_settings={},
llm_call=llm_call,
request=self._request([{"type": "http.disconnect"}]),
)
assert result is model_response
assert not upstream_cancelled.is_set()
async def test_disconnect_cancels_upstream_when_flag_enabled(self, monkeypatch):
upstream_cancelled = asyncio.Event()
async def llm_call():
try:
await asyncio.sleep(5)
return litellm.ModelResponse()
except asyncio.CancelledError:
upstream_cancelled.set()
raise
with pytest.raises(HTTPException) as exc_info:
await self._drive_base_process_llm_request(
monkeypatch,
general_settings={"cancel_on_disconnect": True},
llm_call=llm_call,
request=self._request([{"type": "http.disconnect"}]),
)
assert exc_info.value.status_code == 499
assert upstream_cancelled.is_set()
async def test_499_still_fires_post_call_failure_hook(self):
"""Regression guard: the 499 path must NOT bypass post_call_failure_hook,
which releases max_parallel_requests slots and fires spend/alerting
callbacks (cf. #14457; P1 review finding on #25776/#27146)."""
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
processor = ProxyBaseLLMRequestProcessing(data={})
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
with pytest.raises(ProxyException) as exc_info:
await processor._handle_llm_api_exception(
e=HTTPException(
status_code=499, detail="Client disconnected the request"
),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=proxy_logging_obj,
)
assert exc_info.value.code == "499"
proxy_logging_obj.post_call_failure_hook.assert_awaited_once()
class TestAllmPassthroughRoutePostCallGuardrails:
"""
Regression: non-streaming allm_passthrough_route responses are httpx.Response objects.

View file

@ -42,7 +42,11 @@ _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..",
if _REPO_ROOT not in sys.path:
sys.path.insert(0, _REPO_ROOT)
from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES
from backend.routes.allowlist import (
BACKEND_EXACT_PATHS,
BACKEND_MOUNT_PATHS,
BACKEND_PATH_PREFIXES,
)
from gateway.routes.allowlist import GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES
from litellm.proxy.proxy_server import app
@ -88,3 +92,44 @@ def test_gateway_plus_backend_covers_full_app():
f"Update gateway/routes/allowlist.py or backend/routes/allowlist.py to cover:\n "
+ "\n ".join(sorted(uncovered))
)
def test_backend_mount_paths_defined():
"""BACKEND_MOUNT_PATHS constant must exist and be a frozenset."""
assert isinstance(BACKEND_MOUNT_PATHS, frozenset), \
f"BACKEND_MOUNT_PATHS must be a frozenset, got {type(BACKEND_MOUNT_PATHS)}"
assert len(BACKEND_MOUNT_PATHS) > 0, \
"BACKEND_MOUNT_PATHS must contain at least one Mount path"
def test_swagger_mount_in_backend_allowlist():
"""The /swagger Mount must be in BACKEND_MOUNT_PATHS."""
assert "/swagger" in BACKEND_MOUNT_PATHS, \
"/swagger Mount path must be in BACKEND_MOUNT_PATHS"
def test_backend_keeps_swagger_mount():
"""Verify that Mounts in BACKEND_MOUNT_PATHS are kept on the backend."""
backend_mounts = {
getattr(r, "path")
for r in app.router.routes
if isinstance(r, Mount) and getattr(r, "path", None) in BACKEND_MOUNT_PATHS
}
assert "/swagger" in backend_mounts, \
"/swagger Mount is expected on the proxy app and should be in BACKEND_MOUNT_PATHS"
def test_backend_drops_non_allowlisted_mounts():
"""Verify that Mounts NOT in BACKEND_MOUNT_PATHS would be dropped from backend."""
all_mounts = {
getattr(r, "path")
for r in app.router.routes
if isinstance(r, Mount) and getattr(r, "path", None) is not None
}
non_backend_mounts = all_mounts - BACKEND_MOUNT_PATHS
assert len(non_backend_mounts) > 0, \
"Expected at least one non-backend Mount (e.g., /ui, /_next) to verify filtering logic"
for mount_path in non_backend_mounts:
assert mount_path not in BACKEND_MOUNT_PATHS, \
f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS"

View file

@ -4603,3 +4603,65 @@ def test_apply_overrides_provider_prefix_in_model_skips_router_lookup(
assert data["api_base"] == "https://hotel-eastus.openai.azure.com/"
assert data["api_key"] == "key-hotel-eastus"
router.get_deployment_by_model_group_name.assert_not_called()
def _make_request_mock(path: str, headers: dict) -> MagicMock:
request_mock = MagicMock(spec=Request)
request_mock.url = MagicMock()
request_mock.url.path = path
request_mock.url.__str__.return_value = f"http://localhost{path}"
request_mock.method = "POST"
request_mock.query_params = {}
request_mock.headers = headers
request_mock.client = MagicMock()
request_mock.client.host = "127.0.0.1"
return request_mock
@pytest.mark.asyncio
@pytest.mark.parametrize(
"user_agent, request_drop_params, operator_drop_params, expected_drop_params",
[
("claude-cli/2.0.69 (external, cli)", None, None, True),
("claude-cli/1.0.44 (external, sdk-py)", None, None, True),
("claude-cli/2.0.69 (external, cli)", False, None, False),
("claude-cli/2.0.69 (external, cli)", None, False, None),
("claude-cli/2.0.69 (external, cli)", None, True, None),
("PostmanRuntime/7.53.0", None, None, None),
(None, None, None, None),
],
)
async def test_add_litellm_data_to_request_claude_code_drop_params(
user_agent, request_drop_params, operator_drop_params, expected_drop_params
):
"""Claude Code sends Anthropic-specific params that fail on non-Anthropic
providers, so its user agent must turn on drop_params automatically,
without overriding an explicit caller value, an explicit operator-level
litellm_settings value, or affecting other clients.
"""
headers = {"Content-Type": "application/json"}
if user_agent is not None:
headers["user-agent"] = user_agent
request_mock = _make_request_mock("/v1/messages", headers)
data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
if request_drop_params is not None:
data["drop_params"] = request_drop_params
proxy_config = MagicMock()
proxy_config.config = (
{"litellm_settings": {"drop_params": operator_drop_params}}
if operator_drop_params is not None
else {"litellm_settings": {}}
)
updated = await add_litellm_data_to_request(
data=data,
request=request_mock,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=proxy_config,
general_settings={},
version="test-version",
)
assert updated.get("drop_params") == expected_drop_params

View file

@ -1928,23 +1928,6 @@ async def test_delete_deployment_type_mismatch():
# Create mock ProxyConfig instance
pc = ProxyConfig()
pc.get_config = MagicMock(
return_value={
"model_list": [
{
"model_name": "openai-gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": 12345678},
},
{
"model_name": "openai-gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": 12345679},
},
]
}
)
# Mock llm_router with string IDs (this is the source of the type mismatch)
mock_llm_router = MagicMock()
mock_llm_router.get_model_ids.return_value = [
@ -1963,11 +1946,23 @@ async def test_delete_deployment_type_mismatch():
mock_llm_router.delete_deployment = MagicMock(side_effect=mock_delete_deployment)
# Mock get_config to return empty config (no config models)
async def mock_get_config(config_file_path):
return {}
return {
"model_list": [
{
"model_name": "openai-gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": 12345678},
},
{
"model_name": "openai-gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": 12345679},
},
]
}
pc.get_config = MagicMock(side_effect=mock_get_config)
pc.get_config = AsyncMock(side_effect=mock_get_config)
# Patch the global llm_router
with (
@ -1977,20 +1972,29 @@ async def test_delete_deployment_type_mismatch():
# Call the function under test
deleted_count = await pc._delete_deployment(db_models=[])
# Assertions: Models 12345678 and 12345679 should NOT be deleted
# because they exist in combined_id_list (as integers) even though
# router has them as strings
# The two SHA-hash models have no corresponding entry in combined_id_list
# and must be evicted.
assert (
deleted_count == 2
), f"Expected 2 deletions (SHA-hash models), got {deleted_count}"
assert (
"a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695"
in deleted_ids
)
assert (
"a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3"
in deleted_ids
)
# The function should delete the other 2 models that are not in combined_id_list
assert deleted_count == 0, f"Expected 0 deletions, got {deleted_count}"
# Verify that 12345678 and 12345679 were NOT deleted
assert (
"12345678" not in deleted_ids
), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}"
assert (
"12345679" not in deleted_ids
), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}"
# Models 12345678 and 12345679 exist in the config (as integers); str()
# conversion in _delete_deployment makes them match the router's string IDs,
# so they must NOT be evicted.
assert (
"12345678" not in deleted_ids
), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}"
assert (
"12345679" not in deleted_ids
), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}"
@pytest.mark.asyncio
@ -7937,3 +7941,106 @@ class TestSortModelsByDisplayName:
all_models=models, sort_by="model_name", sort_order="asc"
)
assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"]
class TestDeleteDeploymentSync:
@pytest.mark.asyncio
async def test_delete_deployment_evicts_model_when_all_db_models_deleted(self):
"""
Regression test for #28443.
When all DB models are deleted, _delete_deployment must evict them from
the router. The old code returned 0 early when db_models was empty.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_router = MagicMock()
mock_router.get_model_ids.return_value = ["model-id-to-evict"]
mock_router.delete_deployment.return_value = MagicMock()
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
with patch.object(
proxy_config, "get_config", AsyncMock(return_value={"model_list": []})
):
count = await proxy_config._delete_deployment(db_models=[])
mock_router.delete_deployment.assert_called_once_with(id="model-id-to-evict")
assert count == 1
@pytest.mark.asyncio
async def test_update_llm_router_skips_update_on_db_fetch_failure(self):
"""
When _get_models_from_db returns None (transient DB failure), _update_llm_router
must return early without touching the router.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_router = MagicMock()
with patch("litellm.proxy.proxy_server.llm_router", mock_router):
with patch.object(proxy_config, "get_config", AsyncMock(return_value={})):
await proxy_config._update_llm_router(
new_models=None, proxy_logging_obj=MagicMock()
)
mock_router.delete_deployment.assert_not_called()
mock_router.upsert_deployment.assert_not_called()
@pytest.mark.asyncio
async def test_get_models_from_db_returns_none_on_exception(self):
"""
_get_models_from_db must return None (not []) when the DB raises an exception,
so callers can distinguish a transient failure from a genuinely empty DB.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
side_effect=Exception("DB connection lost")
)
result = await proxy_config._get_models_from_db(prisma_client=mock_prisma)
assert (
result is None
), f"Expected None on DB failure to signal fetch error, got {result!r}"
def test_get_config_list_includes_cancel_on_disconnect(monkeypatch):
"""Follow-up to #30223: the flag must be discoverable via /config/list,
which requires both the ConfigGeneralSettings field and the allowed_args
entry in get_config_list; missing either silently hides it from the UI."""
import types
from unittest.mock import AsyncMock, MagicMock
from fastapi.testclient import TestClient
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.proxy_server import app
mock_prisma = MagicMock()
mock_config_table = MagicMock()
mock_config_table.find_first = AsyncMock(return_value=None)
mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
try:
client = TestClient(app)
resp = client.get("/config/list", params={"config_type": "general_settings"})
assert resp.status_code == 200, resp.text
fields = {item["field_name"]: item for item in resp.json()}
assert "cancel_on_disconnect" in fields
assert fields["cancel_on_disconnect"]["field_type"] == "Boolean"
finally:
app.dependency_overrides.clear()

View file

@ -412,6 +412,20 @@ class TestAnthropicBetaHeadersFiltering:
assert filtered == ["compact-2026-01-12"]
@pytest.mark.parametrize("provider", ["bedrock_converse", "bedrock"])
def test_fine_grained_tool_streaming_forwarded_for_bedrock(self, provider):
"""Bedrock honors fine-grained-tool-streaming-2025-05-14 via
additionalModelRequestFields.anthropic_beta. Stripping it (previously
mapped to null) silently re-enables Anthropic's server-side buffering of
tool-call argument deltas, so streamed tool args arrive in a single
end-of-stream burst instead of incrementally."""
filtered = filter_and_transform_beta_headers(
beta_headers=["fine-grained-tool-streaming-2025-05-14"],
provider=provider,
)
assert filtered == ["fine-grained-tool-streaming-2025-05-14"]
def test_null_value_headers_filtered(self):
"""Test that headers with null values are always filtered out."""
for provider in [

View file

@ -1,5 +1,5 @@
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { render, screen, waitFor } from "@testing-library/react";
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import { Form } from "antd";
import { beforeAll, describe, expect, it, vi } from "vitest";
import { Providers } from "../provider_info_helpers";
@ -215,4 +215,134 @@ describe("ProviderSpecificFields", () => {
expect(baseModelInput).toBeInTheDocument();
});
});
it("sets Azure API version from the API base query parameter", async () => {
const queryClient = createQueryClient();
render(
<QueryClientProvider client={queryClient}>
<Form>
<ProviderSpecificFields selectedProvider={Providers.Azure} />
</Form>
</QueryClientProvider>,
);
const apiBaseInput = await screen.findByPlaceholderText("https://...");
const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview");
fireEvent.change(apiBaseInput, {
target: {
value:
"https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api_version=2024-10-21",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("2024-10-21");
});
});
it("sets Azure API version from the hyphenated API base query parameter", async () => {
const queryClient = createQueryClient();
render(
<QueryClientProvider client={queryClient}>
<Form>
<ProviderSpecificFields selectedProvider={Providers.Azure} />
</Form>
</QueryClientProvider>,
);
const apiBaseInput = await screen.findByPlaceholderText("https://...");
const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview");
fireEvent.change(apiBaseInput, {
target: {
value:
"https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-10-21",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("2024-10-21");
});
});
it("clears an inferred Azure API version when the API base has no version parameter", async () => {
const queryClient = createQueryClient();
render(
<QueryClientProvider client={queryClient}>
<Form>
<ProviderSpecificFields selectedProvider={Providers.Azure} />
</Form>
</QueryClientProvider>,
);
const apiBaseInput = await screen.findByPlaceholderText("https://...");
const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview");
fireEvent.change(apiBaseInput, {
target: {
value:
"https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-10-21",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("2024-10-21");
});
fireEvent.change(apiBaseInput, {
target: {
value: "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("");
});
});
it("preserves a manually edited Azure API version when the API base has no version parameter", async () => {
const queryClient = createQueryClient();
render(
<QueryClientProvider client={queryClient}>
<Form>
<ProviderSpecificFields selectedProvider={Providers.Azure} />
</Form>
</QueryClientProvider>,
);
const apiBaseInput = await screen.findByPlaceholderText("https://...");
const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview");
fireEvent.change(apiBaseInput, {
target: {
value:
"https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-10-21",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("2024-10-21");
});
fireEvent.change(apiVersionInput, {
target: {
value: "2025-01-01-preview",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("2025-01-01-preview");
});
fireEvent.change(apiBaseInput, {
target: {
value: "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions",
},
});
await waitFor(() => {
expect(apiVersionInput).toHaveValue("2025-01-01-preview");
});
});
});

View file

@ -28,6 +28,18 @@ export interface CredentialValues {
value: string;
}
const getApiVersionFromApiBase = (apiBase: string): string | null => {
const queryStartIndex = apiBase.indexOf("?");
if (queryStartIndex === -1) {
return null;
}
const queryString = apiBase.slice(queryStartIndex + 1).split("#")[0];
const searchParams = new URLSearchParams(queryString);
return searchParams.get("api_version") || searchParams.get("api-version");
};
const mapFieldMetadataToUiField = (field: ProviderCredentialFieldMetadata): ProviderCredentialField => {
const type: ProviderCredentialField["type"] =
field.field_type === "password"
@ -167,6 +179,30 @@ const ProviderSpecificFields: React.FC<ProviderSpecificFieldsProps> = ({ selecte
return mapped;
}, [selectedProviderEnum, selectedProvider, providerMetadata]);
const hasApiVersionField = React.useMemo(() => allFields.some((field) => field.key === "api_version"), [allFields]);
const lastInferredApiVersionRef = React.useRef<string | null>(null);
const handleApiBaseChange = React.useCallback(
(event: React.ChangeEvent<HTMLInputElement>) => {
if (!hasApiVersionField) {
return;
}
const apiVersion = getApiVersionFromApiBase(event.target.value);
if (apiVersion) {
lastInferredApiVersionRef.current = apiVersion;
form.setFieldsValue({ api_version: apiVersion });
return;
}
if (form.getFieldValue("api_version") === lastInferredApiVersionRef.current) {
form.setFieldsValue({ api_version: "" });
}
lastInferredApiVersionRef.current = null;
},
[form, hasApiVersionField],
);
const handleUpload = {
name: "file",
accept: ".json",
@ -261,6 +297,7 @@ const ProviderSpecificFields: React.FC<ProviderSpecificFieldsProps> = ({ selecte
placeholder={field.placeholder}
type={field.type === "password" ? "password" : "text"}
defaultValue={field.defaultValue}
onChange={field.key === "api_base" ? handleApiBaseChange : undefined}
/>
)}
</Form.Item>

View file

@ -22062,6 +22062,11 @@ export interface components {
* @description run health checks in background
*/
background_health_checks?: boolean | null;
/**
* Cancel On Disconnect
* @description cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure
*/
cancel_on_disconnect?: boolean | null;
/**
* Completion Model
* @description proxy level default model for all chat completion calls