mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
7d1f68e72a
commit
079c136742
69 changed files with 5055 additions and 633 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
124
litellm/integrations/datadog/datadog_team_handler.py
Normal file
124
litellm/integrations/datadog/datadog_team_handler.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
183
litellm/llms/voyage/embedding/transformation_multimodal.py
Normal file
183
litellm/llms/voyage/embedding/transformation_multimodal.py
Normal 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
|
||||
)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
63
tests/litellm/llms/openai_like/test_empiriolabs_provider.py
Normal file
63
tests/litellm/llms/openai_like/test_empiriolabs_provider.py
Normal 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"
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
@ -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"
|
||||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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 [
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue